{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import os\n",
    "import re\n",
    "import sys\n",
    "import json\n",
    "import pathlib\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.formula.api as sm\n",
    "from collections import Counter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data directory: /shared/nas2/xingyao6/projects/llm-agent/data/outputs\n",
      "Matching glob pattern: `/shared/nas2/xingyao6/projects/llm-agent/data/outputs/**/*results.jsonl`. **46** files found.\n"
     ]
    }
   ],
   "source": [
    "ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.getcwd()))))\n",
    "# Should be your path to the repo `mint`\n",
    "sys.path.insert(0, ROOT_DIR)\n",
    "DATA_DIR = os.path.join(ROOT_DIR, \"data\", \"outputs\")\n",
    "print(f\"Data directory: {DATA_DIR}\")\n",
    "glob_pattern = f\"{DATA_DIR}/**/*results.jsonl\"\n",
    "filepaths = list(set(glob(glob_pattern, recursive=True)))\n",
    "print(f\"Matching glob pattern: `{glob_pattern}`. **{len(filepaths)}** files found.\")\n",
    "\n",
    "\n",
    "def parse_filepath(filepath):\n",
    "    # e.g., gpt-3.5-turbo-0613/F=gpt-3.5-turbo-16k-0613/PHF=no_GT-textual/max5_p2+tool+cd/code_generation/humaneval/results.jsonl\n",
    "    # e.g., gpt-3.5-turbo-0613/F=None/max5_p2+tool+cd/code_generation/humaneval/results.jsonl\n",
    "    splited = filepath.replace(DATA_DIR, \"\").lstrip(\"/\").split(\"/\")\n",
    "    \n",
    "    agent_model_name = splited[0]\n",
    "    feedback_model_name = splited[1].split(\"=\")[1]\n",
    "    if feedback_model_name != \"None\":\n",
    "        feedback_setting = splited[2]\n",
    "    else:\n",
    "        feedback_setting = \"None\"\n",
    "    split = splited[-2]\n",
    "    task_name = splited[-3]\n",
    "    task_type = splited[-4]\n",
    "    exp_setting = splited[-5]\n",
    "    return {\n",
    "        \"agent_model_name\": agent_model_name,\n",
    "        \"feedback_model_name\": feedback_model_name,\n",
    "        \"feedback_setting\": feedback_setting,\n",
    "        \"task_name\": task_name,\n",
    "        \"task_type\": task_type,\n",
    "        \"split\": split,\n",
    "        \"exp_setting\": exp_setting,\n",
    "        \"filepath\": filepath,\n",
    "    }\n",
    "\n",
    "df = pd.DataFrame(list(map(parse_filepath, filepaths)))\n",
    "\n",
    "def load_results(filepath):\n",
    "    results = []\n",
    "    with open(filepath) as f:\n",
    "        for line in f:\n",
    "            try:\n",
    "                results.append(json.loads(line))\n",
    "            except Exception as e:\n",
    "                print(f\"Error loading {filepath}: {e}\\n{line}\")\n",
    "                globals()[\"error_line\"] = line\n",
    "    return pd.DataFrame(results)\n",
    "\n",
    "df[\"results\"] = df.filepath.apply(load_results)\n",
    "\n",
    "\n",
    "def rename_model(model_name):\n",
    "    if \"-hf\" in model_name:\n",
    "        model_name = model_name.rstrip(\"-hf\")\n",
    "    return model_name\n",
    "\n",
    "all_results = []\n",
    "for row in df.itertuples():\n",
    "    row.results[\"agent_model_name\"] = rename_model(row.agent_model_name)\n",
    "    row.results[\"feedback_model_name\"] = rename_model(row.feedback_model_name)\n",
    "    row.results[\"feedback_setting\"] = row.feedback_setting\n",
    "    row.results[\"exp_setting\"] = row.exp_setting\n",
    "    row.results[\"task_name\"] = row.task_name\n",
    "    row.results[\"task_type\"] = row.task_type\n",
    "    all_results.append(row.results)\n",
    "\n",
    "\n",
    "all_results = pd.concat(all_results)\n",
    "def get_stats(row):\n",
    "    state = row[\"state\"]\n",
    "    task = row[\"task\"]\n",
    "    return {\n",
    "        \"task_id\": task[\"task_id\"],\n",
    "        \"n_turns\": len(state[\"history\"]) // 2,\n",
    "        \"success\": state[\"success\"],\n",
    "        \"agent_action_count\": state[\"agent_action_count\"],\n",
    "        \"token_counter\": {'a': Counter(state[\"token_counter\"]), 'b': 1},\n",
    "        \"terminate_reason\": state[\"terminate_reason\"],\n",
    "    }\n",
    "\n",
    "\n",
    "# combine this with the original dataset\n",
    "stats = all_results.apply(get_stats, axis=1, result_type=\"expand\")\n",
    "all_results = pd.concat([all_results, stats], axis=1)\n",
    "all_results[\"token_counter\"] = all_results[\"token_counter\"].apply(lambda x: x[\"a\"])\n",
    "\n",
    "# turn bool to int\n",
    "all_results['success'] = all_results['success'].astype(int)\n",
    "\n",
    "all_results_unfiltered = all_results.copy()\n",
    "\n",
    "# Special handling\n",
    "# make all code_generation from gpt-3.5-turbo-16k-0613 to gpt-3.5-turbo-0613\n",
    "all_results.loc[all_results.agent_model_name == \"gpt-3.5-turbo-16k-0613\", \"agent_model_name\"] = \"gpt-3.5-turbo-0613\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove duplicates in case of weird bugs\n",
    "all_results_no_dup = all_results.drop_duplicates(\n",
    "    subset=[\"task_type\", \"task_name\", \"task_id\", \"agent_model_name\", \"feedback_model_name\", \"feedback_setting\", \"exp_setting\"],\n",
    "    keep=\"first\"\n",
    ")\n",
    "if len(all_results_no_dup) != len(all_results):\n",
    "    print(f\"WARNING: Removed {len(all_results) - len(all_results_no_dup)} duplicated rows.\")\n",
    "    all_results = all_results_no_dup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_7d6e1_row0_col0, #T_7d6e1_row0_col1, #T_7d6e1_row0_col2, #T_7d6e1_row1_col0, #T_7d6e1_row1_col1, #T_7d6e1_row1_col2, #T_7d6e1_row2_col0, #T_7d6e1_row2_col1, #T_7d6e1_row2_col2, #T_7d6e1_row3_col0, #T_7d6e1_row3_col1, #T_7d6e1_row3_col2, #T_7d6e1_row4_col0, #T_7d6e1_row4_col1, #T_7d6e1_row4_col2, #T_7d6e1_row5_col0, #T_7d6e1_row5_col1, #T_7d6e1_row5_col2, #T_7d6e1_row6_col0, #T_7d6e1_row6_col1, #T_7d6e1_row6_col2, #T_7d6e1_row7_col0, #T_7d6e1_row7_col1, #T_7d6e1_row7_col2, #T_7d6e1_row8_col0, #T_7d6e1_row8_col1, #T_7d6e1_row8_col2, #T_7d6e1_row9_col0, #T_7d6e1_row9_col1, #T_7d6e1_row9_col2, #T_7d6e1_row10_col0, #T_7d6e1_row10_col1, #T_7d6e1_row10_col2, #T_7d6e1_row11_col0, #T_7d6e1_row11_col1, #T_7d6e1_row11_col2 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_7d6e1_row0_col3, #T_7d6e1_row1_col3, #T_7d6e1_row2_col3, #T_7d6e1_row3_col3, #T_7d6e1_row4_col3, #T_7d6e1_row5_col3, #T_7d6e1_row6_col3, #T_7d6e1_row7_col3, #T_7d6e1_row8_col3, #T_7d6e1_row9_col3, #T_7d6e1_row10_col3, #T_7d6e1_row11_col3 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_7d6e1\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >agent_model_name</th>\n",
       "      <th id=\"T_7d6e1_level0_col0\" class=\"col_heading level0 col0\" >claude-2</th>\n",
       "      <th id=\"T_7d6e1_level0_col1\" class=\"col_heading level0 col1\" >claude-instant-1</th>\n",
       "      <th id=\"T_7d6e1_level0_col2\" class=\"col_heading level0 col2\" >gpt-3.5-turbo-0613</th>\n",
       "      <th id=\"T_7d6e1_level0_col3\" class=\"col_heading level0 col3\" >gpt-4-0613</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level1\" >feedback_model_name</th>\n",
       "      <th id=\"T_7d6e1_level1_col0\" class=\"col_heading level1 col0\" >None</th>\n",
       "      <th id=\"T_7d6e1_level1_col1\" class=\"col_heading level1 col1\" >None</th>\n",
       "      <th id=\"T_7d6e1_level1_col2\" class=\"col_heading level1 col2\" >None</th>\n",
       "      <th id=\"T_7d6e1_level1_col3\" class=\"col_heading level1 col3\" >None</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level2\" >feedback_setting</th>\n",
       "      <th id=\"T_7d6e1_level2_col0\" class=\"col_heading level2 col0\" >None</th>\n",
       "      <th id=\"T_7d6e1_level2_col1\" class=\"col_heading level2 col1\" >None</th>\n",
       "      <th id=\"T_7d6e1_level2_col2\" class=\"col_heading level2 col2\" >None</th>\n",
       "      <th id=\"T_7d6e1_level2_col3\" class=\"col_heading level2 col3\" >None</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level3\" >exp_setting</th>\n",
       "      <th id=\"T_7d6e1_level3_col0\" class=\"col_heading level3 col0\" >max5_p2+tool+cd</th>\n",
       "      <th id=\"T_7d6e1_level3_col1\" class=\"col_heading level3 col1\" >max5_p2+tool+cd</th>\n",
       "      <th id=\"T_7d6e1_level3_col2\" class=\"col_heading level3 col2\" >max5_p2+tool+cd</th>\n",
       "      <th id=\"T_7d6e1_level3_col3\" class=\"col_heading level3 col3\" >max5_p2+tool+cd</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "      <th class=\"blank col3\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row0\" class=\"row_heading level0 row0\" >APPS</th>\n",
       "      <td id=\"T_7d6e1_row0_col0\" class=\"data row0 col0\" >4439</td>\n",
       "      <td id=\"T_7d6e1_row0_col1\" class=\"data row0 col1\" >4439</td>\n",
       "      <td id=\"T_7d6e1_row0_col2\" class=\"data row0 col2\" >4439</td>\n",
       "      <td id=\"T_7d6e1_row0_col3\" class=\"data row0 col3\" >595</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row1\" class=\"row_heading level0 row1\" >alfworld</th>\n",
       "      <td id=\"T_7d6e1_row1_col0\" class=\"data row1 col0\" >3553</td>\n",
       "      <td id=\"T_7d6e1_row1_col1\" class=\"data row1 col1\" >3553</td>\n",
       "      <td id=\"T_7d6e1_row1_col2\" class=\"data row1 col2\" >3553</td>\n",
       "      <td id=\"T_7d6e1_row1_col3\" class=\"data row1 col3\" >0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row2\" class=\"row_heading level0 row2\" >algebra</th>\n",
       "      <td id=\"T_7d6e1_row2_col0\" class=\"data row2 col0\" >1226</td>\n",
       "      <td id=\"T_7d6e1_row2_col1\" class=\"data row2 col1\" >1226</td>\n",
       "      <td id=\"T_7d6e1_row2_col2\" class=\"data row2 col2\" >1226</td>\n",
       "      <td id=\"T_7d6e1_row2_col3\" class=\"data row2 col3\" >334</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row3\" class=\"row_heading level0 row3\" >counting_and_probability</th>\n",
       "      <td id=\"T_7d6e1_row3_col0\" class=\"data row3 col0\" >602</td>\n",
       "      <td id=\"T_7d6e1_row3_col1\" class=\"data row3 col1\" >602</td>\n",
       "      <td id=\"T_7d6e1_row3_col2\" class=\"data row3 col2\" >602</td>\n",
       "      <td id=\"T_7d6e1_row3_col3\" class=\"data row3 col3\" >311</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row4\" class=\"row_heading level0 row4\" >geometry</th>\n",
       "      <td id=\"T_7d6e1_row4_col0\" class=\"data row4 col0\" >727</td>\n",
       "      <td id=\"T_7d6e1_row4_col1\" class=\"data row4 col1\" >727</td>\n",
       "      <td id=\"T_7d6e1_row4_col2\" class=\"data row4 col2\" >727</td>\n",
       "      <td id=\"T_7d6e1_row4_col3\" class=\"data row4 col3\" >500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row5\" class=\"row_heading level0 row5\" >hotpotqa</th>\n",
       "      <td id=\"T_7d6e1_row5_col0\" class=\"data row5 col0\" >3000</td>\n",
       "      <td id=\"T_7d6e1_row5_col1\" class=\"data row5 col1\" >3000</td>\n",
       "      <td id=\"T_7d6e1_row5_col2\" class=\"data row5 col2\" >3000</td>\n",
       "      <td id=\"T_7d6e1_row5_col3\" class=\"data row5 col3\" >500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row6\" class=\"row_heading level0 row6\" >intermediate_algebra</th>\n",
       "      <td id=\"T_7d6e1_row6_col0\" class=\"data row6 col0\" >1070</td>\n",
       "      <td id=\"T_7d6e1_row6_col1\" class=\"data row6 col1\" >1070</td>\n",
       "      <td id=\"T_7d6e1_row6_col2\" class=\"data row6 col2\" >1070</td>\n",
       "      <td id=\"T_7d6e1_row6_col3\" class=\"data row6 col3\" >500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row7\" class=\"row_heading level0 row7\" >number_theory</th>\n",
       "      <td id=\"T_7d6e1_row7_col0\" class=\"data row7 col0\" >691</td>\n",
       "      <td id=\"T_7d6e1_row7_col1\" class=\"data row7 col1\" >691</td>\n",
       "      <td id=\"T_7d6e1_row7_col2\" class=\"data row7 col2\" >691</td>\n",
       "      <td id=\"T_7d6e1_row7_col3\" class=\"data row7 col3\" >252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row8\" class=\"row_heading level0 row8\" >prealgebra</th>\n",
       "      <td id=\"T_7d6e1_row8_col0\" class=\"data row8 col0\" >750</td>\n",
       "      <td id=\"T_7d6e1_row8_col1\" class=\"data row8 col1\" >750</td>\n",
       "      <td id=\"T_7d6e1_row8_col2\" class=\"data row8 col2\" >750</td>\n",
       "      <td id=\"T_7d6e1_row8_col3\" class=\"data row8 col3\" >172</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row9\" class=\"row_heading level0 row9\" >precalculus</th>\n",
       "      <td id=\"T_7d6e1_row9_col0\" class=\"data row9 col0\" >520</td>\n",
       "      <td id=\"T_7d6e1_row9_col1\" class=\"data row9 col1\" >520</td>\n",
       "      <td id=\"T_7d6e1_row9_col2\" class=\"data row9 col2\" >520</td>\n",
       "      <td id=\"T_7d6e1_row9_col3\" class=\"data row9 col3\" >381</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row10\" class=\"row_heading level0 row10\" >strategyqa</th>\n",
       "      <td id=\"T_7d6e1_row10_col0\" class=\"data row10 col0\" >2290</td>\n",
       "      <td id=\"T_7d6e1_row10_col1\" class=\"data row10 col1\" >2290</td>\n",
       "      <td id=\"T_7d6e1_row10_col2\" class=\"data row10 col2\" >2290</td>\n",
       "      <td id=\"T_7d6e1_row10_col3\" class=\"data row10 col3\" >29</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_7d6e1_level0_row11\" class=\"row_heading level0 row11\" >wiki_table_questions</th>\n",
       "      <td id=\"T_7d6e1_row11_col0\" class=\"data row11 col0\" >3000</td>\n",
       "      <td id=\"T_7d6e1_row11_col1\" class=\"data row11 col1\" >3000</td>\n",
       "      <td id=\"T_7d6e1_row11_col2\" class=\"data row11 col2\" >3000</td>\n",
       "      <td id=\"T_7d6e1_row11_col3\" class=\"data row11 col3\" >0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffa3025a6d0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_f6bcf_row0_col0, #T_f6bcf_row1_col0, #T_f6bcf_row2_col0 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_f6bcf_row3_col0 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_f6bcf\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank\" >&nbsp;</th>\n",
       "      <th class=\"blank\" >&nbsp;</th>\n",
       "      <th class=\"index_name level0\" >exp_setting</th>\n",
       "      <th id=\"T_f6bcf_level0_col0\" class=\"col_heading level0 col0\" >max5_p2+tool+cd</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >agent_model_name</th>\n",
       "      <th class=\"index_name level1\" >feedback_model_name</th>\n",
       "      <th class=\"index_name level2\" >feedback_setting</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_f6bcf_level0_row0\" class=\"row_heading level0 row0\" >claude-2</th>\n",
       "      <th id=\"T_f6bcf_level1_row0\" class=\"row_heading level1 row0\" >None</th>\n",
       "      <th id=\"T_f6bcf_level2_row0\" class=\"row_heading level2 row0\" >None</th>\n",
       "      <td id=\"T_f6bcf_row0_col0\" class=\"data row0 col0\" >21868</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_f6bcf_level0_row1\" class=\"row_heading level0 row1\" >claude-instant-1</th>\n",
       "      <th id=\"T_f6bcf_level1_row1\" class=\"row_heading level1 row1\" >None</th>\n",
       "      <th id=\"T_f6bcf_level2_row1\" class=\"row_heading level2 row1\" >None</th>\n",
       "      <td id=\"T_f6bcf_row1_col0\" class=\"data row1 col0\" >21868</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_f6bcf_level0_row2\" class=\"row_heading level0 row2\" >gpt-3.5-turbo-0613</th>\n",
       "      <th id=\"T_f6bcf_level1_row2\" class=\"row_heading level1 row2\" >None</th>\n",
       "      <th id=\"T_f6bcf_level2_row2\" class=\"row_heading level2 row2\" >None</th>\n",
       "      <td id=\"T_f6bcf_row2_col0\" class=\"data row2 col0\" >21868</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_f6bcf_level0_row3\" class=\"row_heading level0 row3\" >gpt-4-0613</th>\n",
       "      <th id=\"T_f6bcf_level1_row3\" class=\"row_heading level1 row3\" >None</th>\n",
       "      <th id=\"T_f6bcf_level2_row3\" class=\"row_heading level2 row3\" >None</th>\n",
       "      <td id=\"T_f6bcf_row3_col0\" class=\"data row3 col0\" >3574</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffb58693280>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Sanity check of experiments - check whether they are all completed\n",
    "all_results_count = all_results.groupby([\n",
    "    \"agent_model_name\",\n",
    "    \"feedback_model_name\",\n",
    "    \"feedback_setting\",\n",
    "    \"exp_setting\",\n",
    "    # \"task_type\",\n",
    "    \"task_name\",\n",
    "])[\"task_id\"] \\\n",
    ".count().unstack().fillna(0)\\\n",
    "# .sum(axis=1).unstack().fillna(0).astype(int)\n",
    "\n",
    "display(all_results_count.T.astype(int).style.background_gradient(cmap='Blues', axis=1))\n",
    "display(all_results_count.sum(axis=1).unstack().fillna(0).astype(int).style.background_gradient(cmap='Blues', axis=0))\n",
    "\n",
    "# separate the results from gpt4\n",
    "gpt4_results = all_results[all_results.agent_model_name.str.startswith(\"gpt-4\")]\n",
    "all_results = all_results[~all_results.agent_model_name.str.startswith(\"gpt-4\")]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>n_turns</th>\n",
       "      <th>success</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>464</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>463</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>419</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>300</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>305</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>498</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>499</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3574 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     n_turns  success\n",
       "464        1        1\n",
       "463        1        1\n",
       "419        1        1\n",
       "300        1        1\n",
       "305        1        1\n",
       "..       ...      ...\n",
       "6          5        0\n",
       "5          5        0\n",
       "498        5        0\n",
       "499        5        0\n",
       "0          5        0\n",
       "\n",
       "[3574 rows x 2 columns]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.concat(\n",
    "    [gpt4_results, gpt4_results[\"token_counter\"].apply(pd.Series)],\n",
    "    axis=1\n",
    ")[[\"n_turns\", \"success\"]].sort_values(\"n_turns\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>n_turns</th>\n",
       "      <th>success</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>464</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>463</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>419</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>300</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>305</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>498</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>499</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3574 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     n_turns  success\n",
       "464        1        1\n",
       "463        1        1\n",
       "419        1        1\n",
       "300        1        1\n",
       "305        1        1\n",
       "..       ...      ...\n",
       "6          5        0\n",
       "5          5        0\n",
       "498        5        0\n",
       "499        5        0\n",
       "0          5        0\n",
       "\n",
       "[3574 rows x 2 columns]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gpt4_results[[\"n_turns\", \"success\"]].sort_values(\"n_turns\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr:last-of-type th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <th colspan=\"8\" halign=\"left\">n_turns</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <th>count</th>\n",
       "      <th>mean</th>\n",
       "      <th>std</th>\n",
       "      <th>min</th>\n",
       "      <th>25%</th>\n",
       "      <th>50%</th>\n",
       "      <th>75%</th>\n",
       "      <th>max</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>success</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2535.0</td>\n",
       "      <td>4.344773</td>\n",
       "      <td>1.055423</td>\n",
       "      <td>2.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1039.0</td>\n",
       "      <td>2.666987</td>\n",
       "      <td>1.079132</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        n_turns                                             \n",
       "          count      mean       std  min  25%  50%  75%  max\n",
       "success                                                     \n",
       "0        2535.0  4.344773  1.055423  2.0  4.0  5.0  5.0  5.0\n",
       "1        1039.0  2.666987  1.079132  1.0  2.0  2.0  3.0  5.0"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.concat(\n",
    "    [gpt4_results, gpt4_results[\"token_counter\"].apply(pd.Series)],\n",
    "    axis=1\n",
    ")[[\"n_turns\", \"success\"]].groupby(\"success\")[[\"n_turns\"]].describe()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "423"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(gpt4_results.query(\"n_turns > 2 and success == 1\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "APPS                        4439\n",
       "alfworld                    3553\n",
       "algebra                     1226\n",
       "counting_and_probability     602\n",
       "geometry                     727\n",
       "hotpotqa                    3000\n",
       "intermediate_algebra        1070\n",
       "number_theory                691\n",
       "prealgebra                   750\n",
       "precalculus                  520\n",
       "strategyqa                  2290\n",
       "wiki_table_questions        3000\n",
       "dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before filtering: 65604\n",
      "After filtering: 65604\n",
      "agent_model_name    feedback_model_name  feedback_setting  exp_setting    \n",
      "claude-2            None                 None              max5_p2+tool+cd    True\n",
      "claude-instant-1    None                 None              max5_p2+tool+cd    True\n",
      "gpt-3.5-turbo-0613  None                 None              max5_p2+tool+cd    True\n",
      "dtype: bool\n"
     ]
    }
   ],
   "source": [
    "# Filter out experiments that are not completed\n",
    "\n",
    "# find all index that are not [136, 134, 320]\n",
    "TASK_COUNT = {\n",
    "    \"prealgebra\": 750,\n",
    "    \"number_theory\": 691,\n",
    "    \"algebra\": 1226,\n",
    "    \"precalculus\": 520,\n",
    "    \"hotpotqa\": 3000,\n",
    "    \"counting_and_probability\": 602,\n",
    "    \"strategyqa\": 2290,\n",
    "    \"intermediate_algebra\": 1070,\n",
    "    \"geometry\": 727,\n",
    "    \"APPS\": 4439, # 2156,\n",
    "    \"alfworld\": 3553,\n",
    "    \"wiki_table_questions\": 3000\n",
    "}\n",
    "\n",
    "# GLOBAL_MAX = all_results_count.max()\n",
    "# assert (GLOBAL_MAX == pd.Series([136, 134, 316], index=[\"code_generation\", \"decision_making\", \"reasoning\"])).all()\n",
    "TASK_COUNT_ROW = pd.Series(TASK_COUNT).sort_index()\n",
    "display(TASK_COUNT_ROW)\n",
    "def _exp_completed(row):\n",
    "    assert len(row) == len(TASK_COUNT_ROW), f\"row: {row}, TASK_COUNT_ROW: {TASK_COUNT_ROW}\"\n",
    "    # sort by index\n",
    "    row = row.sort_index()\n",
    "    return (row == TASK_COUNT_ROW).all()\n",
    "\n",
    "completed_exp = all_results_count.apply(_exp_completed, axis=1)\n",
    "# select only completed exp\n",
    "completed_exp = completed_exp.drop(completed_exp[completed_exp == False].index)#.reset_index().drop(columns=[0])\n",
    "# display(completed_exp.to_frame().style.background_gradient(cmap='Blues', axis=0))\n",
    "\n",
    "completed_exp_lst = set(map(tuple, completed_exp.reset_index().drop(columns=[0]).to_numpy().tolist()))\n",
    "# agent_model_name\tfeedback_model_name\tfeedback_setting\texp_setting\n",
    "# completed_exp_lst\n",
    "_completed_mask = all_results.apply(lambda row: (row[\"agent_model_name\"], row[\"feedback_model_name\"], row[\"feedback_setting\"], row[\"exp_setting\"]) in completed_exp_lst, axis=1)\n",
    "print(f\"Before filtering: {len(all_results)}\")\n",
    "not_completed = all_results[~_completed_mask]\n",
    "completed_results = all_results[_completed_mask]\n",
    "completed_results_w_stats = pd.concat([\n",
    "    completed_results,\n",
    "    completed_results[\"agent_action_count\"].apply(pd.Series)\n",
    "], axis=1)\n",
    "print(f\"After filtering: {len(completed_results)}\")\n",
    "print(completed_exp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt4_completed_results_w_stats = pd.concat([\n",
    "    gpt4_results,\n",
    "    gpt4_results[\"agent_action_count\"].apply(pd.Series)\n",
    "], axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['claude-2', 'claude-instant-1', 'gpt-3.5-turbo-0613']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FILTERED_MODELS =[\n",
    "    # TODO(user): fill in models you want to filter out\n",
    "]\n",
    "all_results = all_results[~all_results[\"agent_model_name\"].isin(FILTERED_MODELS)]\n",
    "sorted(list(all_results[\"agent_model_name\"].unique()))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cost Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "** No feedback **\n",
      "Total rounds of feedback: 57521\n",
      "Cost per 1k tokens:\n",
      "{'prompt_tokens': 0.0015, 'completion_tokens': 0.002, 'token_count': 0, 'feedback_prompt_tokens': 0.03, 'feedback_completion_tokens': 0.04}\n",
      "Total Tokens\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "n_tasks                  21868\n",
       "prompt_tokens        131130326\n",
       "completion_tokens     10727062\n",
       "total_tokens         141857388\n",
       "dtype: int64"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of tokens per example\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>n_tasks</th>\n",
       "      <th>prompt_tokens</th>\n",
       "      <th>completion_tokens</th>\n",
       "      <th>total_tokens</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>task_type</th>\n",
       "      <th>task_name</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>code_generation</th>\n",
       "      <th>APPS</th>\n",
       "      <td>4439</td>\n",
       "      <td>8695.328452</td>\n",
       "      <td>799.622212</td>\n",
       "      <td>9494.950665</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>decision_making</th>\n",
       "      <th>alfworld</th>\n",
       "      <td>3553</td>\n",
       "      <td>7361.853645</td>\n",
       "      <td>508.993527</td>\n",
       "      <td>7870.847171</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"9\" valign=\"top\">reasoning</th>\n",
       "      <th>algebra</th>\n",
       "      <td>1226</td>\n",
       "      <td>4333.333605</td>\n",
       "      <td>413.539152</td>\n",
       "      <td>4746.872757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>counting_and_probability</th>\n",
       "      <td>602</td>\n",
       "      <td>4911.335548</td>\n",
       "      <td>544.951827</td>\n",
       "      <td>5456.287375</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>geometry</th>\n",
       "      <td>727</td>\n",
       "      <td>6246.950481</td>\n",
       "      <td>707.078404</td>\n",
       "      <td>6954.028886</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>hotpotqa</th>\n",
       "      <td>3000</td>\n",
       "      <td>5531.042333</td>\n",
       "      <td>275.697000</td>\n",
       "      <td>5806.739333</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>intermediate_algebra</th>\n",
       "      <td>1070</td>\n",
       "      <td>6260.781308</td>\n",
       "      <td>765.708411</td>\n",
       "      <td>7026.489720</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>number_theory</th>\n",
       "      <td>691</td>\n",
       "      <td>4423.749638</td>\n",
       "      <td>408.120116</td>\n",
       "      <td>4831.869754</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>prealgebra</th>\n",
       "      <td>750</td>\n",
       "      <td>3800.381333</td>\n",
       "      <td>328.548000</td>\n",
       "      <td>4128.929333</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>precalculus</th>\n",
       "      <td>520</td>\n",
       "      <td>6542.832692</td>\n",
       "      <td>831.525000</td>\n",
       "      <td>7374.357692</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>strategyqa</th>\n",
       "      <td>2290</td>\n",
       "      <td>3636.600437</td>\n",
       "      <td>215.204803</td>\n",
       "      <td>3851.805240</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>tabular</th>\n",
       "      <th>wiki_table_questions</th>\n",
       "      <td>3000</td>\n",
       "      <td>4211.641667</td>\n",
       "      <td>306.648667</td>\n",
       "      <td>4518.290333</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                          n_tasks  prompt_tokens  \\\n",
       "task_type       task_name                                          \n",
       "code_generation APPS                         4439    8695.328452   \n",
       "decision_making alfworld                     3553    7361.853645   \n",
       "reasoning       algebra                      1226    4333.333605   \n",
       "                counting_and_probability      602    4911.335548   \n",
       "                geometry                      727    6246.950481   \n",
       "                hotpotqa                     3000    5531.042333   \n",
       "                intermediate_algebra         1070    6260.781308   \n",
       "                number_theory                 691    4423.749638   \n",
       "                prealgebra                    750    3800.381333   \n",
       "                precalculus                   520    6542.832692   \n",
       "                strategyqa                   2290    3636.600437   \n",
       "tabular         wiki_table_questions         3000    4211.641667   \n",
       "\n",
       "                                          completion_tokens  total_tokens  \n",
       "task_type       task_name                                                  \n",
       "code_generation APPS                             799.622212   9494.950665  \n",
       "decision_making alfworld                         508.993527   7870.847171  \n",
       "reasoning       algebra                          413.539152   4746.872757  \n",
       "                counting_and_probability         544.951827   5456.287375  \n",
       "                geometry                         707.078404   6954.028886  \n",
       "                hotpotqa                         275.697000   5806.739333  \n",
       "                intermediate_algebra             765.708411   7026.489720  \n",
       "                number_theory                    408.120116   4831.869754  \n",
       "                prealgebra                       328.548000   4128.929333  \n",
       "                precalculus                      831.525000   7374.357692  \n",
       "                strategyqa                       215.204803   3851.805240  \n",
       "tabular         wiki_table_questions             306.648667   4518.290333  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_2a9c5_row0_col2 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_2a9c5_row1_col2 {\n",
       "  background-color: #2d7dbb;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_2a9c5_row2_col2 {\n",
       "  background-color: #d6e6f4;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_2a9c5_row3_col2 {\n",
       "  background-color: #b8d5ea;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_2a9c5_row4_col2 {\n",
       "  background-color: #57a0ce;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_2a9c5_row5_col2 {\n",
       "  background-color: #aacfe5;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_2a9c5_row6_col2 {\n",
       "  background-color: #529dcc;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_2a9c5_row7_col2 {\n",
       "  background-color: #d4e4f4;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_2a9c5_row8_col2 {\n",
       "  background-color: #edf4fc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_2a9c5_row9_col2 {\n",
       "  background-color: #3e8ec4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_2a9c5_row10_col2 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_2a9c5_row11_col2 {\n",
       "  background-color: #dfecf7;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_2a9c5\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank\" >&nbsp;</th>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_2a9c5_level0_col0\" class=\"col_heading level0 col0\" >prompt_tokens</th>\n",
       "      <th id=\"T_2a9c5_level0_col1\" class=\"col_heading level0 col1\" >completion_tokens</th>\n",
       "      <th id=\"T_2a9c5_level0_col2\" class=\"col_heading level0 col2\" >USD_per_example</th>\n",
       "      <th id=\"T_2a9c5_level0_col3\" class=\"col_heading level0 col3\" >n_tasks</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_type</th>\n",
       "      <th class=\"index_name level1\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "      <th class=\"blank col3\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level0_row0\" class=\"row_heading level0 row0\" >code_generation</th>\n",
       "      <th id=\"T_2a9c5_level1_row0\" class=\"row_heading level1 row0\" >APPS</th>\n",
       "      <td id=\"T_2a9c5_row0_col0\" class=\"data row0 col0\" >0.013043</td>\n",
       "      <td id=\"T_2a9c5_row0_col1\" class=\"data row0 col1\" >0.001599</td>\n",
       "      <td id=\"T_2a9c5_row0_col2\" class=\"data row0 col2\" >0.014642</td>\n",
       "      <td id=\"T_2a9c5_row0_col3\" class=\"data row0 col3\" >4439</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level0_row1\" class=\"row_heading level0 row1\" >decision_making</th>\n",
       "      <th id=\"T_2a9c5_level1_row1\" class=\"row_heading level1 row1\" >alfworld</th>\n",
       "      <td id=\"T_2a9c5_row1_col0\" class=\"data row1 col0\" >0.011043</td>\n",
       "      <td id=\"T_2a9c5_row1_col1\" class=\"data row1 col1\" >0.001018</td>\n",
       "      <td id=\"T_2a9c5_row1_col2\" class=\"data row1 col2\" >0.012061</td>\n",
       "      <td id=\"T_2a9c5_row1_col3\" class=\"data row1 col3\" >3553</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level0_row2\" class=\"row_heading level0 row2\" rowspan=\"9\">reasoning</th>\n",
       "      <th id=\"T_2a9c5_level1_row2\" class=\"row_heading level1 row2\" >algebra</th>\n",
       "      <td id=\"T_2a9c5_row2_col0\" class=\"data row2 col0\" >0.006500</td>\n",
       "      <td id=\"T_2a9c5_row2_col1\" class=\"data row2 col1\" >0.000827</td>\n",
       "      <td id=\"T_2a9c5_row2_col2\" class=\"data row2 col2\" >0.007327</td>\n",
       "      <td id=\"T_2a9c5_row2_col3\" class=\"data row2 col3\" >1226</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row3\" class=\"row_heading level1 row3\" >counting_and_probability</th>\n",
       "      <td id=\"T_2a9c5_row3_col0\" class=\"data row3 col0\" >0.007367</td>\n",
       "      <td id=\"T_2a9c5_row3_col1\" class=\"data row3 col1\" >0.001090</td>\n",
       "      <td id=\"T_2a9c5_row3_col2\" class=\"data row3 col2\" >0.008457</td>\n",
       "      <td id=\"T_2a9c5_row3_col3\" class=\"data row3 col3\" >602</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row4\" class=\"row_heading level1 row4\" >geometry</th>\n",
       "      <td id=\"T_2a9c5_row4_col0\" class=\"data row4 col0\" >0.009370</td>\n",
       "      <td id=\"T_2a9c5_row4_col1\" class=\"data row4 col1\" >0.001414</td>\n",
       "      <td id=\"T_2a9c5_row4_col2\" class=\"data row4 col2\" >0.010785</td>\n",
       "      <td id=\"T_2a9c5_row4_col3\" class=\"data row4 col3\" >727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row5\" class=\"row_heading level1 row5\" >hotpotqa</th>\n",
       "      <td id=\"T_2a9c5_row5_col0\" class=\"data row5 col0\" >0.008297</td>\n",
       "      <td id=\"T_2a9c5_row5_col1\" class=\"data row5 col1\" >0.000551</td>\n",
       "      <td id=\"T_2a9c5_row5_col2\" class=\"data row5 col2\" >0.008848</td>\n",
       "      <td id=\"T_2a9c5_row5_col3\" class=\"data row5 col3\" >3000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row6\" class=\"row_heading level1 row6\" >intermediate_algebra</th>\n",
       "      <td id=\"T_2a9c5_row6_col0\" class=\"data row6 col0\" >0.009391</td>\n",
       "      <td id=\"T_2a9c5_row6_col1\" class=\"data row6 col1\" >0.001531</td>\n",
       "      <td id=\"T_2a9c5_row6_col2\" class=\"data row6 col2\" >0.010923</td>\n",
       "      <td id=\"T_2a9c5_row6_col3\" class=\"data row6 col3\" >1070</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row7\" class=\"row_heading level1 row7\" >number_theory</th>\n",
       "      <td id=\"T_2a9c5_row7_col0\" class=\"data row7 col0\" >0.006636</td>\n",
       "      <td id=\"T_2a9c5_row7_col1\" class=\"data row7 col1\" >0.000816</td>\n",
       "      <td id=\"T_2a9c5_row7_col2\" class=\"data row7 col2\" >0.007452</td>\n",
       "      <td id=\"T_2a9c5_row7_col3\" class=\"data row7 col3\" >691</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row8\" class=\"row_heading level1 row8\" >prealgebra</th>\n",
       "      <td id=\"T_2a9c5_row8_col0\" class=\"data row8 col0\" >0.005701</td>\n",
       "      <td id=\"T_2a9c5_row8_col1\" class=\"data row8 col1\" >0.000657</td>\n",
       "      <td id=\"T_2a9c5_row8_col2\" class=\"data row8 col2\" >0.006358</td>\n",
       "      <td id=\"T_2a9c5_row8_col3\" class=\"data row8 col3\" >750</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row9\" class=\"row_heading level1 row9\" >precalculus</th>\n",
       "      <td id=\"T_2a9c5_row9_col0\" class=\"data row9 col0\" >0.009814</td>\n",
       "      <td id=\"T_2a9c5_row9_col1\" class=\"data row9 col1\" >0.001663</td>\n",
       "      <td id=\"T_2a9c5_row9_col2\" class=\"data row9 col2\" >0.011477</td>\n",
       "      <td id=\"T_2a9c5_row9_col3\" class=\"data row9 col3\" >520</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level1_row10\" class=\"row_heading level1 row10\" >strategyqa</th>\n",
       "      <td id=\"T_2a9c5_row10_col0\" class=\"data row10 col0\" >0.005455</td>\n",
       "      <td id=\"T_2a9c5_row10_col1\" class=\"data row10 col1\" >0.000430</td>\n",
       "      <td id=\"T_2a9c5_row10_col2\" class=\"data row10 col2\" >0.005885</td>\n",
       "      <td id=\"T_2a9c5_row10_col3\" class=\"data row10 col3\" >2290</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_2a9c5_level0_row11\" class=\"row_heading level0 row11\" >tabular</th>\n",
       "      <th id=\"T_2a9c5_level1_row11\" class=\"row_heading level1 row11\" >wiki_table_questions</th>\n",
       "      <td id=\"T_2a9c5_row11_col0\" class=\"data row11 col0\" >0.006317</td>\n",
       "      <td id=\"T_2a9c5_row11_col1\" class=\"data row11 col1\" >0.000613</td>\n",
       "      <td id=\"T_2a9c5_row11_col2\" class=\"data row11 col2\" >0.006931</td>\n",
       "      <td id=\"T_2a9c5_row11_col3\" class=\"data row11 col3\" >3000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7fde7c00ac40>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total cost: $218.15\n"
     ]
    }
   ],
   "source": [
    "def get_token_count(df, cost_per_1k_token):\n",
    "    n_round_of_feedback = (df[\"n_turns\"] - 1).sum()\n",
    "    print(f\"Total rounds of feedback: {n_round_of_feedback}\")\n",
    "\n",
    "    print(f\"Cost per 1k tokens:\\n{cost_per_1k_token}\")\n",
    "    token_counter = df.groupby([\"task_type\", \"task_name\"])[\"token_counter\"].apply(lambda x: (sum(x, Counter()), len(x)))\n",
    "    token_counter = pd.DataFrame(token_counter.tolist(), index=token_counter.index, columns=[\"token_counter\", \"n_tasks\"])\n",
    "    # expand the token_counter\n",
    "    token_counter = pd.concat([token_counter.drop(columns=[\"token_counter\"]), token_counter[\"token_counter\"].apply(pd.Series)], axis=1)\n",
    "    # normalize by n_tasks\n",
    "    n_tasks = token_counter[\"n_tasks\"]\n",
    "    print(\"Total Tokens\")\n",
    "    display(token_counter.sum(axis=0))\n",
    "    token_counter = token_counter.div(token_counter.n_tasks, axis=0)#.drop(columns=[\"n_tasks\"])\n",
    "    token_counter[\"n_tasks\"] = n_tasks\n",
    "    print(f\"Number of tokens per example\")\n",
    "\n",
    "    display(token_counter)\n",
    "\n",
    "    costs_per_ex = token_counter.drop(columns=[\"n_tasks\"])\n",
    "    if \"total_tokens\" in costs_per_ex.columns:\n",
    "        costs_per_ex = costs_per_ex.drop(columns=[\"total_tokens\"])\n",
    "    if \"feedback_total_tokens\" in costs_per_ex.columns:\n",
    "        costs_per_ex = costs_per_ex.drop(columns=[\"feedback_total_tokens\"])\n",
    "\n",
    "    costs_per_ex = costs_per_ex.div(1000).apply(lambda x: x * cost_per_1k_token[x.name])\n",
    "    costs_per_ex[\"USD_per_example\"] = costs_per_ex.sum(axis=1)\n",
    "    costs_per_ex[\"n_tasks\"] = token_counter[\"n_tasks\"]\n",
    "    # display(costs_per_ex)\n",
    "    # styler gradient\n",
    "    display(costs_per_ex.style.background_gradient(cmap=\"Blues\", axis=None, subset=[\"USD_per_example\"]))\n",
    "\n",
    "    # Total cost\n",
    "    total_cost = (costs_per_ex[\"USD_per_example\"] * costs_per_ex[\"n_tasks\"]).sum()\n",
    "    print(f\"Total cost: ${total_cost:.2f}\")\n",
    "\n",
    "print(\"** No feedback **\")\n",
    "GPT35_COST_PER_1K_TOKEN = {\n",
    "    # 3.5-turbo\n",
    "    \"prompt_tokens\": 0.0015,\n",
    "    \"completion_tokens\": 0.002,\n",
    "    # chat-bison-001\n",
    "    \"token_count\": 0,\n",
    "    # 4\n",
    "    # \"prompt_tokens\": 0.03,\n",
    "    # \"completion_tokens\": 0.04,\n",
    "    # 3.5-turbo-16k\n",
    "    # \"feedback_prompt_tokens\": 0.003,\n",
    "    # \"feedback_completion_tokens\": 0.004,\n",
    "    # 4\n",
    "    \"feedback_prompt_tokens\": 0.03,\n",
    "    \"feedback_completion_tokens\": 0.04,\n",
    "}\n",
    "get_token_count(all_results.query(\"agent_model_name == 'gpt-3.5-turbo-0613' and feedback_setting == 'None'\"), GPT35_COST_PER_1K_TOKEN)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_failed = completed_results_w_stats[\"success\"].apply(lambda x: not bool(x))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_9b7be_row0_col0 {\n",
       "  background-color: #60a7d2;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row0_col1, #T_9b7be_row0_col4, #T_9b7be_row0_col6, #T_9b7be_row0_col7, #T_9b7be_row10_col0, #T_9b7be_row10_col2, #T_9b7be_row10_col3, #T_9b7be_row10_col5, #T_9b7be_row10_col8 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row0_col2 {\n",
       "  background-color: #d2e3f3;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row0_col3 {\n",
       "  background-color: #e4eff9;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row0_col5, #T_9b7be_row4_col6, #T_9b7be_row4_col8, #T_9b7be_row9_col0, #T_9b7be_row9_col1, #T_9b7be_row9_col2, #T_9b7be_row9_col3, #T_9b7be_row9_col4, #T_9b7be_row9_col6, #T_9b7be_row9_col7 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row0_col8 {\n",
       "  background-color: #4292c6;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row1_col0 {\n",
       "  background-color: #c2d9ee;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row1_col1, #T_9b7be_row1_col4, #T_9b7be_row1_col7 {\n",
       "  background-color: #1c6ab0;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row1_col2, #T_9b7be_row7_col0 {\n",
       "  background-color: #e8f1fa;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row1_col3 {\n",
       "  background-color: #08509b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row1_col5 {\n",
       "  background-color: #5fa6d1;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row1_col6 {\n",
       "  background-color: #084184;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row1_col8 {\n",
       "  background-color: #1f6eb3;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row2_col0 {\n",
       "  background-color: #ccdff1;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row2_col1, #T_9b7be_row2_col4, #T_9b7be_row2_col7 {\n",
       "  background-color: #d3e4f3;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row2_col2 {\n",
       "  background-color: #9dcae1;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row2_col3, #T_9b7be_row7_col2 {\n",
       "  background-color: #c7dbef;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row2_col5, #T_9b7be_row11_col5 {\n",
       "  background-color: #79b5d9;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row2_col6 {\n",
       "  background-color: #bfd8ed;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row2_col8 {\n",
       "  background-color: #3080bd;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row3_col0, #T_9b7be_row6_col0, #T_9b7be_row7_col1, #T_9b7be_row7_col4, #T_9b7be_row7_col7 {\n",
       "  background-color: #eef5fc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row3_col1, #T_9b7be_row3_col4, #T_9b7be_row3_col7 {\n",
       "  background-color: #f3f8fe;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row3_col2 {\n",
       "  background-color: #d3e3f3;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row3_col3 {\n",
       "  background-color: #eff6fc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row3_col5 {\n",
       "  background-color: #cfe1f2;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row3_col6, #T_9b7be_row9_col5 {\n",
       "  background-color: #ecf4fb;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row3_col8 {\n",
       "  background-color: #a6cee4;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row4_col0 {\n",
       "  background-color: #f2f7fd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row4_col1, #T_9b7be_row4_col4, #T_9b7be_row4_col5, #T_9b7be_row4_col7, #T_9b7be_row6_col3 {\n",
       "  background-color: #edf4fc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row4_col2, #T_9b7be_row9_col8 {\n",
       "  background-color: #eaf3fb;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row4_col3 {\n",
       "  background-color: #f5fafe;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row5_col0 {\n",
       "  background-color: #3585bf;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row5_col1, #T_9b7be_row5_col4, #T_9b7be_row5_col7, #T_9b7be_row11_col1, #T_9b7be_row11_col4, #T_9b7be_row11_col7 {\n",
       "  background-color: #3f8fc5;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row5_col2 {\n",
       "  background-color: #7cb7da;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row5_col3 {\n",
       "  background-color: #2b7bba;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row5_col5 {\n",
       "  background-color: #64a9d3;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row5_col6 {\n",
       "  background-color: #57a0ce;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row5_col8 {\n",
       "  background-color: #69add5;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row6_col1, #T_9b7be_row6_col4, #T_9b7be_row6_col7 {\n",
       "  background-color: #dce9f6;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row6_col2 {\n",
       "  background-color: #f0f6fd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row6_col5 {\n",
       "  background-color: #e1edf8;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row6_col6 {\n",
       "  background-color: #eaf2fb;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row6_col8 {\n",
       "  background-color: #d8e7f5;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row7_col3 {\n",
       "  background-color: #e7f1fa;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row7_col5 {\n",
       "  background-color: #b4d3e9;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row7_col6 {\n",
       "  background-color: #ddeaf7;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row7_col8 {\n",
       "  background-color: #3888c1;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row8_col0 {\n",
       "  background-color: #d5e5f4;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row8_col1, #T_9b7be_row8_col4, #T_9b7be_row8_col7 {\n",
       "  background-color: #ebf3fb;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row8_col2 {\n",
       "  background-color: #5da5d1;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row8_col3 {\n",
       "  background-color: #d9e8f5;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row8_col5 {\n",
       "  background-color: #6caed6;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row8_col6 {\n",
       "  background-color: #d7e6f5;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row8_col8 {\n",
       "  background-color: #2474b7;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row10_col1, #T_9b7be_row10_col4, #T_9b7be_row10_col7 {\n",
       "  background-color: #7fb9da;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row10_col6 {\n",
       "  background-color: #2272b6;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row11_col0 {\n",
       "  background-color: #75b4d8;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row11_col2 {\n",
       "  background-color: #b9d6ea;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_9b7be_row11_col3 {\n",
       "  background-color: #3d8dc4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row11_col6 {\n",
       "  background-color: #3484bf;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_9b7be_row11_col8 {\n",
       "  background-color: #3e8ec4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_9b7be\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >agent_model_name</th>\n",
       "      <th id=\"T_9b7be_level0_col0\" class=\"col_heading level0 col0\" colspan=\"3\">claude-2</th>\n",
       "      <th id=\"T_9b7be_level0_col3\" class=\"col_heading level0 col3\" colspan=\"3\">claude-instant-1</th>\n",
       "      <th id=\"T_9b7be_level0_col6\" class=\"col_heading level0 col6\" colspan=\"3\">gpt-3.5-turbo-0613</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"blank level1\" >&nbsp;</th>\n",
       "      <th id=\"T_9b7be_level1_col0\" class=\"col_heading level1 col0\" >n_success_tasks</th>\n",
       "      <th id=\"T_9b7be_level1_col1\" class=\"col_heading level1 col1\" >n_tasks</th>\n",
       "      <th id=\"T_9b7be_level1_col2\" class=\"col_heading level1 col2\" >success_rate</th>\n",
       "      <th id=\"T_9b7be_level1_col3\" class=\"col_heading level1 col3\" >n_success_tasks</th>\n",
       "      <th id=\"T_9b7be_level1_col4\" class=\"col_heading level1 col4\" >n_tasks</th>\n",
       "      <th id=\"T_9b7be_level1_col5\" class=\"col_heading level1 col5\" >success_rate</th>\n",
       "      <th id=\"T_9b7be_level1_col6\" class=\"col_heading level1 col6\" >n_success_tasks</th>\n",
       "      <th id=\"T_9b7be_level1_col7\" class=\"col_heading level1 col7\" >n_tasks</th>\n",
       "      <th id=\"T_9b7be_level1_col8\" class=\"col_heading level1 col8\" >success_rate</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "      <th class=\"blank col3\" >&nbsp;</th>\n",
       "      <th class=\"blank col4\" >&nbsp;</th>\n",
       "      <th class=\"blank col5\" >&nbsp;</th>\n",
       "      <th class=\"blank col6\" >&nbsp;</th>\n",
       "      <th class=\"blank col7\" >&nbsp;</th>\n",
       "      <th class=\"blank col8\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row0\" class=\"row_heading level0 row0\" >APPS</th>\n",
       "      <td id=\"T_9b7be_row0_col0\" class=\"data row0 col0\" >1187</td>\n",
       "      <td id=\"T_9b7be_row0_col1\" class=\"data row0 col1\" >4439</td>\n",
       "      <td id=\"T_9b7be_row0_col2\" class=\"data row0 col2\" >0.267403</td>\n",
       "      <td id=\"T_9b7be_row0_col3\" class=\"data row0 col3\" >247</td>\n",
       "      <td id=\"T_9b7be_row0_col4\" class=\"data row0 col4\" >4439</td>\n",
       "      <td id=\"T_9b7be_row0_col5\" class=\"data row0 col5\" >0.055643</td>\n",
       "      <td id=\"T_9b7be_row0_col6\" class=\"data row0 col6\" >2109</td>\n",
       "      <td id=\"T_9b7be_row0_col7\" class=\"data row0 col7\" >4439</td>\n",
       "      <td id=\"T_9b7be_row0_col8\" class=\"data row0 col8\" >0.475107</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row1\" class=\"row_heading level0 row1\" >alfworld</th>\n",
       "      <td id=\"T_9b7be_row1_col0\" class=\"data row1 col0\" >617</td>\n",
       "      <td id=\"T_9b7be_row1_col1\" class=\"data row1 col1\" >3553</td>\n",
       "      <td id=\"T_9b7be_row1_col2\" class=\"data row1 col2\" >0.173656</td>\n",
       "      <td id=\"T_9b7be_row1_col3\" class=\"data row1 col3\" >1826</td>\n",
       "      <td id=\"T_9b7be_row1_col4\" class=\"data row1 col4\" >3553</td>\n",
       "      <td id=\"T_9b7be_row1_col5\" class=\"data row1 col5\" >0.513932</td>\n",
       "      <td id=\"T_9b7be_row1_col6\" class=\"data row1 col6\" >1975</td>\n",
       "      <td id=\"T_9b7be_row1_col7\" class=\"data row1 col7\" >3553</td>\n",
       "      <td id=\"T_9b7be_row1_col8\" class=\"data row1 col8\" >0.555868</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row2\" class=\"row_heading level0 row2\" >algebra</th>\n",
       "      <td id=\"T_9b7be_row2_col0\" class=\"data row2 col0\" >521</td>\n",
       "      <td id=\"T_9b7be_row2_col1\" class=\"data row2 col1\" >1226</td>\n",
       "      <td id=\"T_9b7be_row2_col2\" class=\"data row2 col2\" >0.424959</td>\n",
       "      <td id=\"T_9b7be_row2_col3\" class=\"data row2 col3\" >557</td>\n",
       "      <td id=\"T_9b7be_row2_col4\" class=\"data row2 col4\" >1226</td>\n",
       "      <td id=\"T_9b7be_row2_col5\" class=\"data row2 col5\" >0.454323</td>\n",
       "      <td id=\"T_9b7be_row2_col6\" class=\"data row2 col6\" >633</td>\n",
       "      <td id=\"T_9b7be_row2_col7\" class=\"data row2 col7\" >1226</td>\n",
       "      <td id=\"T_9b7be_row2_col8\" class=\"data row2 col8\" >0.516313</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row3\" class=\"row_heading level0 row3\" >counting_and_probability</th>\n",
       "      <td id=\"T_9b7be_row3_col0\" class=\"data row3 col0\" >160</td>\n",
       "      <td id=\"T_9b7be_row3_col1\" class=\"data row3 col1\" >602</td>\n",
       "      <td id=\"T_9b7be_row3_col2\" class=\"data row3 col2\" >0.265781</td>\n",
       "      <td id=\"T_9b7be_row3_col3\" class=\"data row3 col3\" >138</td>\n",
       "      <td id=\"T_9b7be_row3_col4\" class=\"data row3 col4\" >602</td>\n",
       "      <td id=\"T_9b7be_row3_col5\" class=\"data row3 col5\" >0.229236</td>\n",
       "      <td id=\"T_9b7be_row3_col6\" class=\"data row3 col6\" >190</td>\n",
       "      <td id=\"T_9b7be_row3_col7\" class=\"data row3 col7\" >602</td>\n",
       "      <td id=\"T_9b7be_row3_col8\" class=\"data row3 col8\" >0.315615</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row4\" class=\"row_heading level0 row4\" >geometry</th>\n",
       "      <td id=\"T_9b7be_row4_col0\" class=\"data row4 col0\" >119</td>\n",
       "      <td id=\"T_9b7be_row4_col1\" class=\"data row4 col1\" >727</td>\n",
       "      <td id=\"T_9b7be_row4_col2\" class=\"data row4 col2\" >0.163686</td>\n",
       "      <td id=\"T_9b7be_row4_col3\" class=\"data row4 col3\" >73</td>\n",
       "      <td id=\"T_9b7be_row4_col4\" class=\"data row4 col4\" >727</td>\n",
       "      <td id=\"T_9b7be_row4_col5\" class=\"data row4 col5\" >0.100413</td>\n",
       "      <td id=\"T_9b7be_row4_col6\" class=\"data row4 col6\" >81</td>\n",
       "      <td id=\"T_9b7be_row4_col7\" class=\"data row4 col7\" >727</td>\n",
       "      <td id=\"T_9b7be_row4_col8\" class=\"data row4 col8\" >0.111417</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row5\" class=\"row_heading level0 row5\" >hotpotqa</th>\n",
       "      <td id=\"T_9b7be_row5_col0\" class=\"data row5 col0\" >1482</td>\n",
       "      <td id=\"T_9b7be_row5_col1\" class=\"data row5 col1\" >3000</td>\n",
       "      <td id=\"T_9b7be_row5_col2\" class=\"data row5 col2\" >0.494000</td>\n",
       "      <td id=\"T_9b7be_row5_col3\" class=\"data row5 col3\" >1497</td>\n",
       "      <td id=\"T_9b7be_row5_col4\" class=\"data row5 col4\" >3000</td>\n",
       "      <td id=\"T_9b7be_row5_col5\" class=\"data row5 col5\" >0.499000</td>\n",
       "      <td id=\"T_9b7be_row5_col6\" class=\"data row5 col6\" >1218</td>\n",
       "      <td id=\"T_9b7be_row5_col7\" class=\"data row5 col7\" >3000</td>\n",
       "      <td id=\"T_9b7be_row5_col8\" class=\"data row5 col8\" >0.406000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row6\" class=\"row_heading level0 row6\" >intermediate_algebra</th>\n",
       "      <td id=\"T_9b7be_row6_col0\" class=\"data row6 col0\" >151</td>\n",
       "      <td id=\"T_9b7be_row6_col1\" class=\"data row6 col1\" >1070</td>\n",
       "      <td id=\"T_9b7be_row6_col2\" class=\"data row6 col2\" >0.141121</td>\n",
       "      <td id=\"T_9b7be_row6_col3\" class=\"data row6 col3\" >160</td>\n",
       "      <td id=\"T_9b7be_row6_col4\" class=\"data row6 col4\" >1070</td>\n",
       "      <td id=\"T_9b7be_row6_col5\" class=\"data row6 col5\" >0.149533</td>\n",
       "      <td id=\"T_9b7be_row6_col6\" class=\"data row6 col6\" >219</td>\n",
       "      <td id=\"T_9b7be_row6_col7\" class=\"data row6 col7\" >1070</td>\n",
       "      <td id=\"T_9b7be_row6_col8\" class=\"data row6 col8\" >0.204673</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row7\" class=\"row_heading level0 row7\" >number_theory</th>\n",
       "      <td id=\"T_9b7be_row7_col0\" class=\"data row7 col0\" >220</td>\n",
       "      <td id=\"T_9b7be_row7_col1\" class=\"data row7 col1\" >691</td>\n",
       "      <td id=\"T_9b7be_row7_col2\" class=\"data row7 col2\" >0.318379</td>\n",
       "      <td id=\"T_9b7be_row7_col3\" class=\"data row7 col3\" >218</td>\n",
       "      <td id=\"T_9b7be_row7_col4\" class=\"data row7 col4\" >691</td>\n",
       "      <td id=\"T_9b7be_row7_col5\" class=\"data row7 col5\" >0.315485</td>\n",
       "      <td id=\"T_9b7be_row7_col6\" class=\"data row7 col6\" >344</td>\n",
       "      <td id=\"T_9b7be_row7_col7\" class=\"data row7 col7\" >691</td>\n",
       "      <td id=\"T_9b7be_row7_col8\" class=\"data row7 col8\" >0.497829</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row8\" class=\"row_heading level0 row8\" >prealgebra</th>\n",
       "      <td id=\"T_9b7be_row8_col0\" class=\"data row8 col0\" >423</td>\n",
       "      <td id=\"T_9b7be_row8_col1\" class=\"data row8 col1\" >750</td>\n",
       "      <td id=\"T_9b7be_row8_col2\" class=\"data row8 col2\" >0.564000</td>\n",
       "      <td id=\"T_9b7be_row8_col3\" class=\"data row8 col3\" >360</td>\n",
       "      <td id=\"T_9b7be_row8_col4\" class=\"data row8 col4\" >750</td>\n",
       "      <td id=\"T_9b7be_row8_col5\" class=\"data row8 col5\" >0.480000</td>\n",
       "      <td id=\"T_9b7be_row8_col6\" class=\"data row8 col6\" >406</td>\n",
       "      <td id=\"T_9b7be_row8_col7\" class=\"data row8 col7\" >750</td>\n",
       "      <td id=\"T_9b7be_row8_col8\" class=\"data row8 col8\" >0.541333</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row9\" class=\"row_heading level0 row9\" >precalculus</th>\n",
       "      <td id=\"T_9b7be_row9_col0\" class=\"data row9 col0\" >57</td>\n",
       "      <td id=\"T_9b7be_row9_col1\" class=\"data row9 col1\" >520</td>\n",
       "      <td id=\"T_9b7be_row9_col2\" class=\"data row9 col2\" >0.109615</td>\n",
       "      <td id=\"T_9b7be_row9_col3\" class=\"data row9 col3\" >54</td>\n",
       "      <td id=\"T_9b7be_row9_col4\" class=\"data row9 col4\" >520</td>\n",
       "      <td id=\"T_9b7be_row9_col5\" class=\"data row9 col5\" >0.103846</td>\n",
       "      <td id=\"T_9b7be_row9_col6\" class=\"data row9 col6\" >78</td>\n",
       "      <td id=\"T_9b7be_row9_col7\" class=\"data row9 col7\" >520</td>\n",
       "      <td id=\"T_9b7be_row9_col8\" class=\"data row9 col8\" >0.150000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row10\" class=\"row_heading level0 row10\" >strategyqa</th>\n",
       "      <td id=\"T_9b7be_row10_col0\" class=\"data row10 col0\" >2169</td>\n",
       "      <td id=\"T_9b7be_row10_col1\" class=\"data row10 col1\" >2290</td>\n",
       "      <td id=\"T_9b7be_row10_col2\" class=\"data row10 col2\" >0.947162</td>\n",
       "      <td id=\"T_9b7be_row10_col3\" class=\"data row10 col3\" >2077</td>\n",
       "      <td id=\"T_9b7be_row10_col4\" class=\"data row10 col4\" >2290</td>\n",
       "      <td id=\"T_9b7be_row10_col5\" class=\"data row10 col5\" >0.906987</td>\n",
       "      <td id=\"T_9b7be_row10_col6\" class=\"data row10 col6\" >1592</td>\n",
       "      <td id=\"T_9b7be_row10_col7\" class=\"data row10 col7\" >2290</td>\n",
       "      <td id=\"T_9b7be_row10_col8\" class=\"data row10 col8\" >0.695197</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_9b7be_level0_row11\" class=\"row_heading level0 row11\" >wiki_table_questions</th>\n",
       "      <td id=\"T_9b7be_row11_col0\" class=\"data row11 col0\" >1062</td>\n",
       "      <td id=\"T_9b7be_row11_col1\" class=\"data row11 col1\" >3000</td>\n",
       "      <td id=\"T_9b7be_row11_col2\" class=\"data row11 col2\" >0.354000</td>\n",
       "      <td id=\"T_9b7be_row11_col3\" class=\"data row11 col3\" >1356</td>\n",
       "      <td id=\"T_9b7be_row11_col4\" class=\"data row11 col4\" >3000</td>\n",
       "      <td id=\"T_9b7be_row11_col5\" class=\"data row11 col5\" >0.452000</td>\n",
       "      <td id=\"T_9b7be_row11_col6\" class=\"data row11 col6\" >1451</td>\n",
       "      <td id=\"T_9b7be_row11_col7\" class=\"data row11 col7\" >3000</td>\n",
       "      <td id=\"T_9b7be_row11_col8\" class=\"data row11 col8\" >0.483667</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffa111fa3d0>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_groupby = [\"agent_model_name\", \"task_name\"]\n",
    "pd.concat([\n",
    "    completed_results_w_stats.groupby(_groupby)[\"success\"].count().rename(\"n_tasks\"),\n",
    "    completed_results_w_stats.groupby(_groupby)[\"success\"].sum().rename(\"n_success_tasks\"),\n",
    "    completed_results_w_stats.groupby(_groupby)[\"success\"].mean().rename(\"success_rate\"),\n",
    "    # completed_results_w_stats.groupby(_groupby)[\"success\"].std().rename(\"std\"),\n",
    "], axis=1).unstack(0).swaplevel(0, 1, axis=1).sort_index(axis=1).style.background_gradient(cmap=\"Blues\", axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean number of actions (Success)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_6ceb9_row0_col0 {\n",
       "  background-color: #09529d;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row0_col1 {\n",
       "  background-color: #eef5fc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row0_col2 {\n",
       "  background-color: #3686c0;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row1_col0, #T_6ceb9_row1_col2, #T_6ceb9_row4_col1 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row1_col1, #T_6ceb9_row5_col2, #T_6ceb9_row9_col0 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row2_col0 {\n",
       "  background-color: #084488;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row2_col1 {\n",
       "  background-color: #dbe9f6;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row2_col2 {\n",
       "  background-color: #7db8da;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row3_col0 {\n",
       "  background-color: #083a7a;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row3_col1 {\n",
       "  background-color: #deebf7;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row3_col2 {\n",
       "  background-color: #4d99ca;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row4_col0 {\n",
       "  background-color: #083c7d;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row4_col2 {\n",
       "  background-color: #3d8dc4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row5_col0 {\n",
       "  background-color: #08519c;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row5_col1 {\n",
       "  background-color: #b3d3e8;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row6_col0 {\n",
       "  background-color: #083573;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row6_col1 {\n",
       "  background-color: #abd0e6;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row6_col2 {\n",
       "  background-color: #56a0ce;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row7_col0, #T_6ceb9_row10_col0 {\n",
       "  background-color: #084082;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row7_col1 {\n",
       "  background-color: #b9d6ea;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row7_col2 {\n",
       "  background-color: #ccdff1;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row8_col0 {\n",
       "  background-color: #08478d;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row8_col1 {\n",
       "  background-color: #f0f6fd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row8_col2 {\n",
       "  background-color: #8cc0dd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row9_col1 {\n",
       "  background-color: #c7dcef;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row9_col2 {\n",
       "  background-color: #2272b6;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row10_col1 {\n",
       "  background-color: #d9e8f5;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row10_col2 {\n",
       "  background-color: #2474b7;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row11_col0 {\n",
       "  background-color: #08458a;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_6ceb9_row11_col1 {\n",
       "  background-color: #b2d2e8;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_6ceb9_row11_col2 {\n",
       "  background-color: #edf4fc;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_6ceb9\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_6ceb9_level0_col0\" class=\"col_heading level0 col0\" >propose_solution</th>\n",
       "      <th id=\"T_6ceb9_level0_col1\" class=\"col_heading level0 col1\" >use_tool</th>\n",
       "      <th id=\"T_6ceb9_level0_col2\" class=\"col_heading level0 col2\" >invalid_action</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row0\" class=\"row_heading level0 row0\" >APPS</th>\n",
       "      <td id=\"T_6ceb9_row0_col0\" class=\"data row0 col0\" >1.104714</td>\n",
       "      <td id=\"T_6ceb9_row0_col1\" class=\"data row0 col1\" >0.714366</td>\n",
       "      <td id=\"T_6ceb9_row0_col2\" class=\"data row0 col2\" >0.076207</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row1\" class=\"row_heading level0 row1\" >alfworld</th>\n",
       "      <td id=\"T_6ceb9_row1_col0\" class=\"data row1 col0\" >0.057266</td>\n",
       "      <td id=\"T_6ceb9_row1_col1\" class=\"data row1 col1\" >2.578769</td>\n",
       "      <td id=\"T_6ceb9_row1_col2\" class=\"data row1 col2\" >0.000453</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row2\" class=\"row_heading level0 row2\" >algebra</th>\n",
       "      <td id=\"T_6ceb9_row2_col0\" class=\"data row2 col0\" >1.170660</td>\n",
       "      <td id=\"T_6ceb9_row2_col1\" class=\"data row2 col1\" >0.897721</td>\n",
       "      <td id=\"T_6ceb9_row2_col2\" class=\"data row2 col2\" >0.052016</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row3\" class=\"row_heading level0 row3\" >counting_and_probability</th>\n",
       "      <td id=\"T_6ceb9_row3_col0\" class=\"data row3 col0\" >1.213115</td>\n",
       "      <td id=\"T_6ceb9_row3_col1\" class=\"data row3 col1\" >0.864754</td>\n",
       "      <td id=\"T_6ceb9_row3_col2\" class=\"data row3 col2\" >0.067623</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row4\" class=\"row_heading level0 row4\" >geometry</th>\n",
       "      <td id=\"T_6ceb9_row4_col0\" class=\"data row4 col0\" >1.205128</td>\n",
       "      <td id=\"T_6ceb9_row4_col1\" class=\"data row4 col1\" >0.619048</td>\n",
       "      <td id=\"T_6ceb9_row4_col2\" class=\"data row4 col2\" >0.073260</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row5\" class=\"row_heading level0 row5\" >hotpotqa</th>\n",
       "      <td id=\"T_6ceb9_row5_col0\" class=\"data row5 col0\" >1.111032</td>\n",
       "      <td id=\"T_6ceb9_row5_col1\" class=\"data row5 col1\" >1.226114</td>\n",
       "      <td id=\"T_6ceb9_row5_col2\" class=\"data row5 col2\" >0.113653</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row6\" class=\"row_heading level0 row6\" >intermediate_algebra</th>\n",
       "      <td id=\"T_6ceb9_row6_col0\" class=\"data row6 col0\" >1.237736</td>\n",
       "      <td id=\"T_6ceb9_row6_col1\" class=\"data row6 col1\" >1.273585</td>\n",
       "      <td id=\"T_6ceb9_row6_col2\" class=\"data row6 col2\" >0.064151</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row7\" class=\"row_heading level0 row7\" >number_theory</th>\n",
       "      <td id=\"T_6ceb9_row7_col0\" class=\"data row7 col0\" >1.193095</td>\n",
       "      <td id=\"T_6ceb9_row7_col1\" class=\"data row7 col1\" >1.190537</td>\n",
       "      <td id=\"T_6ceb9_row7_col2\" class=\"data row7 col2\" >0.025575</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row8\" class=\"row_heading level0 row8\" >prealgebra</th>\n",
       "      <td id=\"T_6ceb9_row8_col0\" class=\"data row8 col0\" >1.160639</td>\n",
       "      <td id=\"T_6ceb9_row8_col1\" class=\"data row8 col1\" >0.690496</td>\n",
       "      <td id=\"T_6ceb9_row8_col2\" class=\"data row8 col2\" >0.047939</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row9\" class=\"row_heading level0 row9\" >precalculus</th>\n",
       "      <td id=\"T_6ceb9_row9_col0\" class=\"data row9 col0\" >1.264550</td>\n",
       "      <td id=\"T_6ceb9_row9_col1\" class=\"data row9 col1\" >1.095238</td>\n",
       "      <td id=\"T_6ceb9_row9_col2\" class=\"data row9 col2\" >0.084656</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row10\" class=\"row_heading level0 row10\" >strategyqa</th>\n",
       "      <td id=\"T_6ceb9_row10_col0\" class=\"data row10 col0\" >1.189620</td>\n",
       "      <td id=\"T_6ceb9_row10_col1\" class=\"data row10 col1\" >0.917437</td>\n",
       "      <td id=\"T_6ceb9_row10_col2\" class=\"data row10 col2\" >0.083762</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_6ceb9_level0_row11\" class=\"row_heading level0 row11\" >wiki_table_questions</th>\n",
       "      <td id=\"T_6ceb9_row11_col0\" class=\"data row11 col0\" >1.168002</td>\n",
       "      <td id=\"T_6ceb9_row11_col1\" class=\"data row11 col1\" >1.235720</td>\n",
       "      <td id=\"T_6ceb9_row11_col2\" class=\"data row11 col2\" >0.006203</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffa22b112b0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Mean number of actions (Success)\")\n",
    "completed_results_w_stats[~is_failed].groupby(\"task_name\")[[\"propose_solution\", \"use_tool\", \"invalid_action\"]].mean().style.background_gradient(cmap=\"Blues\", axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean number of actions (Failed)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_b93a0_row0_col0, #T_b93a0_row1_col1, #T_b93a0_row10_col2 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row0_col1 {\n",
       "  background-color: #deebf7;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row0_col2 {\n",
       "  background-color: #c4daee;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row1_col0, #T_b93a0_row1_col2, #T_b93a0_row8_col1 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row2_col0 {\n",
       "  background-color: #2070b4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row2_col1 {\n",
       "  background-color: #dae8f6;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row2_col2 {\n",
       "  background-color: #a8cee4;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row3_col0, #T_b93a0_row8_col0 {\n",
       "  background-color: #083370;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row3_col1 {\n",
       "  background-color: #f0f6fd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row3_col2 {\n",
       "  background-color: #a4cce3;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row4_col0 {\n",
       "  background-color: #3585bf;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row4_col1 {\n",
       "  background-color: #d7e6f5;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row4_col2, #T_b93a0_row11_col1 {\n",
       "  background-color: #65aad4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row5_col0 {\n",
       "  background-color: #b4d3e9;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row5_col1 {\n",
       "  background-color: #539ecd;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row5_col2 {\n",
       "  background-color: #1562a9;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row6_col0 {\n",
       "  background-color: #82bbdb;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row6_col1 {\n",
       "  background-color: #b0d2e7;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row6_col2 {\n",
       "  background-color: #66abd4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row7_col0 {\n",
       "  background-color: #084a91;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row7_col1 {\n",
       "  background-color: #d2e3f3;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row7_col2 {\n",
       "  background-color: #d1e2f3;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row8_col2 {\n",
       "  background-color: #abd0e6;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row9_col0 {\n",
       "  background-color: #72b2d8;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row9_col1 {\n",
       "  background-color: #b9d6ea;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row9_col2 {\n",
       "  background-color: #2d7dbb;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row10_col0 {\n",
       "  background-color: #8dc1dd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b93a0_row10_col1 {\n",
       "  background-color: #4292c6;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row11_col0 {\n",
       "  background-color: #08316d;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b93a0_row11_col2 {\n",
       "  background-color: #eaf3fb;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_b93a0\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_b93a0_level0_col0\" class=\"col_heading level0 col0\" >propose_solution</th>\n",
       "      <th id=\"T_b93a0_level0_col1\" class=\"col_heading level0 col1\" >use_tool</th>\n",
       "      <th id=\"T_b93a0_level0_col2\" class=\"col_heading level0 col2\" >invalid_action</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row0\" class=\"row_heading level0 row0\" >APPS</th>\n",
       "      <td id=\"T_b93a0_row0_col0\" class=\"data row0 col0\" >1.900757</td>\n",
       "      <td id=\"T_b93a0_row0_col1\" class=\"data row0 col1\" >1.243299</td>\n",
       "      <td id=\"T_b93a0_row0_col2\" class=\"data row0 col2\" >0.121649</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row1\" class=\"row_heading level0 row1\" >alfworld</th>\n",
       "      <td id=\"T_b93a0_row1_col0\" class=\"data row1 col0\" >1.542541</td>\n",
       "      <td id=\"T_b93a0_row1_col1\" class=\"data row1 col1\" >2.675212</td>\n",
       "      <td id=\"T_b93a0_row1_col2\" class=\"data row1 col2\" >0.004967</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row2\" class=\"row_heading level0 row2\" >algebra</th>\n",
       "      <td id=\"T_b93a0_row2_col0\" class=\"data row2 col0\" >1.811388</td>\n",
       "      <td id=\"T_b93a0_row2_col1\" class=\"data row2 col1\" >1.277580</td>\n",
       "      <td id=\"T_b93a0_row2_col2\" class=\"data row2 col2\" >0.162684</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row3\" class=\"row_heading level0 row3\" >counting_and_probability</th>\n",
       "      <td id=\"T_b93a0_row3_col0\" class=\"data row3 col0\" >1.895296</td>\n",
       "      <td id=\"T_b93a0_row3_col1\" class=\"data row3 col1\" >1.100910</td>\n",
       "      <td id=\"T_b93a0_row3_col2\" class=\"data row3 col2\" >0.167678</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row4\" class=\"row_heading level0 row4\" >geometry</th>\n",
       "      <td id=\"T_b93a0_row4_col0\" class=\"data row4 col0\" >1.783543</td>\n",
       "      <td id=\"T_b93a0_row4_col1\" class=\"data row4 col1\" >1.302411</td>\n",
       "      <td id=\"T_b93a0_row4_col2\" class=\"data row4 col2\" >0.240566</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row5\" class=\"row_heading level0 row5\" >hotpotqa</th>\n",
       "      <td id=\"T_b93a0_row5_col0\" class=\"data row5 col0\" >1.652925</td>\n",
       "      <td id=\"T_b93a0_row5_col1\" class=\"data row5 col1\" >1.974807</td>\n",
       "      <td id=\"T_b93a0_row5_col2\" class=\"data row5 col2\" >0.372684</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row6\" class=\"row_heading level0 row6\" >intermediate_algebra</th>\n",
       "      <td id=\"T_b93a0_row6_col0\" class=\"data row6 col0\" >1.701866</td>\n",
       "      <td id=\"T_b93a0_row6_col1\" class=\"data row6 col1\" >1.562687</td>\n",
       "      <td id=\"T_b93a0_row6_col2\" class=\"data row6 col2\" >0.238806</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row7\" class=\"row_heading level0 row7\" >number_theory</th>\n",
       "      <td id=\"T_b93a0_row7_col0\" class=\"data row7 col0\" >1.865221</td>\n",
       "      <td id=\"T_b93a0_row7_col1\" class=\"data row7 col1\" >1.349342</td>\n",
       "      <td id=\"T_b93a0_row7_col2\" class=\"data row7 col2\" >0.093726</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row8\" class=\"row_heading level0 row8\" >prealgebra</th>\n",
       "      <td id=\"T_b93a0_row8_col0\" class=\"data row8 col0\" >1.895382</td>\n",
       "      <td id=\"T_b93a0_row8_col1\" class=\"data row8 col1\" >1.038643</td>\n",
       "      <td id=\"T_b93a0_row8_col2\" class=\"data row8 col2\" >0.157399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row9\" class=\"row_heading level0 row9\" >precalculus</th>\n",
       "      <td id=\"T_b93a0_row9_col0\" class=\"data row9 col0\" >1.715536</td>\n",
       "      <td id=\"T_b93a0_row9_col1\" class=\"data row9 col1\" >1.512035</td>\n",
       "      <td id=\"T_b93a0_row9_col2\" class=\"data row9 col2\" >0.326039</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row10\" class=\"row_heading level0 row10\" >strategyqa</th>\n",
       "      <td id=\"T_b93a0_row10_col0\" class=\"data row10 col0\" >1.691860</td>\n",
       "      <td id=\"T_b93a0_row10_col1\" class=\"data row10 col1\" >2.056202</td>\n",
       "      <td id=\"T_b93a0_row10_col2\" class=\"data row10 col2\" >0.460271</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b93a0_level0_row11\" class=\"row_heading level0 row11\" >wiki_table_questions</th>\n",
       "      <td id=\"T_b93a0_row11_col0\" class=\"data row11 col0\" >1.898850</td>\n",
       "      <td id=\"T_b93a0_row11_col1\" class=\"data row11 col1\" >1.884233</td>\n",
       "      <td id=\"T_b93a0_row11_col2\" class=\"data row11 col2\" >0.034496</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffa30244820>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Mean number of actions (Failed)\")\n",
    "completed_results_w_stats[is_failed].groupby(\"task_name\")[[\"propose_solution\", \"use_tool\", \"invalid_action\"]].mean().style.background_gradient(cmap=\"Blues\", axis=0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Easy Tier - These model can complete with <= 2 steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of easy instances: 18447 / 65604 (28.12%)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of unique easy instances: 11277 / 21868 (51.57%)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_24806_row0_col0 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_24806_row1_col0 {\n",
       "  background-color: #3686c0;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_24806_row2_col0 {\n",
       "  background-color: #3f8fc5;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_24806_row3_col0 {\n",
       "  background-color: #4b98ca;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_24806_row4_col0 {\n",
       "  background-color: #60a7d2;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_24806_row5_col0 {\n",
       "  background-color: #bfd8ed;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_24806_row6_col0 {\n",
       "  background-color: #cfe1f2;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_24806_row7_col0 {\n",
       "  background-color: #e4eff9;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_24806_row8_col0 {\n",
       "  background-color: #ebf3fb;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_24806_row9_col0 {\n",
       "  background-color: #eef5fc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_24806_row10_col0 {\n",
       "  background-color: #f2f8fd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_24806_row11_col0 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_24806\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_24806_level0_col0\" class=\"col_heading level0 col0\" >n_tasks</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row0\" class=\"row_heading level0 row0\" >strategyqa</th>\n",
       "      <td id=\"T_24806_row0_col0\" class=\"data row0 col0\" >4243</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row1\" class=\"row_heading level0 row1\" >APPS</th>\n",
       "      <td id=\"T_24806_row1_col0\" class=\"data row1 col0\" >2876</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row2\" class=\"row_heading level0 row2\" >wiki_table_questions</th>\n",
       "      <td id=\"T_24806_row2_col0\" class=\"data row2 col0\" >2740</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row3\" class=\"row_heading level0 row3\" >hotpotqa</th>\n",
       "      <td id=\"T_24806_row3_col0\" class=\"data row3 col0\" >2567</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row4\" class=\"row_heading level0 row4\" >alfworld</th>\n",
       "      <td id=\"T_24806_row4_col0\" class=\"data row4 col0\" >2319</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row5\" class=\"row_heading level0 row5\" >algebra</th>\n",
       "      <td id=\"T_24806_row5_col0\" class=\"data row5 col0\" >1244</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row6\" class=\"row_heading level0 row6\" >prealgebra</th>\n",
       "      <td id=\"T_24806_row6_col0\" class=\"data row6 col0\" >955</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row7\" class=\"row_heading level0 row7\" >number_theory</th>\n",
       "      <td id=\"T_24806_row7_col0\" class=\"data row7 col0\" >512</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row8\" class=\"row_heading level0 row8\" >counting_and_probability</th>\n",
       "      <td id=\"T_24806_row8_col0\" class=\"data row8 col0\" >368</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row9\" class=\"row_heading level0 row9\" >intermediate_algebra</th>\n",
       "      <td id=\"T_24806_row9_col0\" class=\"data row9 col0\" >294</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row10\" class=\"row_heading level0 row10\" >geometry</th>\n",
       "      <td id=\"T_24806_row10_col0\" class=\"data row10 col0\" >213</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_24806_level0_row11\" class=\"row_heading level0 row11\" >precalculus</th>\n",
       "      <td id=\"T_24806_row11_col0\" class=\"data row11 col0\" >116</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffa1166a130>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "easy_instances = completed_results_w_stats[~is_failed].query(\"n_turns <= 2\")\n",
    "print(f\"Number of easy instances: {len(easy_instances)} / {len(completed_results_w_stats)} ({len(easy_instances) / len(completed_results_w_stats):.2%})\")\n",
    "\n",
    "all_unique_instances_task_name_id_pairs = set(completed_results_w_stats[[\"task_type\", \"task_name\", \"task_id\"]].apply(tuple, axis=1))\n",
    "easy_instances_task_name_id_pairs = set(easy_instances[[\"task_type\", \"task_name\", \"task_id\"]].apply(tuple, axis=1))\n",
    "print(f\"Number of unique easy instances: {len(easy_instances_task_name_id_pairs)} / {len(all_unique_instances_task_name_id_pairs)} ({len(easy_instances_task_name_id_pairs) / len(all_unique_instances_task_name_id_pairs):.2%})\")\n",
    "\n",
    "easy_instances\\\n",
    "    .groupby(\"task_name\")[\"success\"].count().rename(\"n_tasks\")\\\n",
    "    .sort_values(ascending=False).to_frame().style.background_gradient(cmap=\"Blues\", axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 2876 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/apps.jsonl\n",
      "Saved 2319 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/alfworld.jsonl\n",
      "Saved 1244 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/algebra.jsonl\n",
      "Saved 368 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/counting_and_probability.jsonl\n",
      "Saved 213 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/geometry.jsonl\n",
      "Saved 2567 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/hotpotqa.jsonl\n",
      "Saved 294 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/intermediate_algebra.jsonl\n",
      "Saved 512 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/number_theory.jsonl\n",
      "Saved 955 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/prealgebra.jsonl\n",
      "Saved 116 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/precalculus.jsonl\n",
      "Saved 4243 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/strategyqa.jsonl\n",
      "Saved 2740 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/easy_instances/wiki_table_questions.jsonl\n"
     ]
    }
   ],
   "source": [
    "# output dir data/datasets/easy_instances\n",
    "def save_instances(df, output_dir):\n",
    "    pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
    "    # Save the easy instances by task_name\n",
    "    for task_name, cur_df in df.groupby(\"task_name\"):\n",
    "        cur_output_dir = os.path.join(output_dir, f\"{task_name.lower()}.jsonl\")\n",
    "        cur_df.to_json(\n",
    "            cur_output_dir,\n",
    "            orient=\"records\",\n",
    "            lines=True\n",
    "        )\n",
    "        print(f\"Saved {len(cur_df)} instances to {cur_output_dir}\")\n",
    "\n",
    "save_instances(easy_instances, os.path.join(ROOT_DIR, \"data\", \"trajectories\", \"easy_instances\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Medium Tier - These model can success with 5 steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of medium instances: 8580 / 65604 (13.08%)\n",
      "Number of unique medium instances: 6941 / 21868 (31.74%)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_b3b1d_row0_col0 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b3b1d_row1_col0 {\n",
       "  background-color: #1c6bb0;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b3b1d_row2_col0 {\n",
       "  background-color: #2070b4;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b3b1d_row3_col0 {\n",
       "  background-color: #63a8d3;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_b3b1d_row4_col0 {\n",
       "  background-color: #b7d4ea;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row5_col0 {\n",
       "  background-color: #d0e1f2;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row6_col0 {\n",
       "  background-color: #e3eef8;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row7_col0 {\n",
       "  background-color: #e6f0f9;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row8_col0 {\n",
       "  background-color: #e7f0fa;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row9_col0 {\n",
       "  background-color: #f2f7fd;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row10_col0 {\n",
       "  background-color: #f6faff;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_b3b1d_row11_col0 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_b3b1d\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_b3b1d_level0_col0\" class=\"col_heading level0 col0\" >n_tasks</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row0\" class=\"row_heading level0 row0\" >alfworld</th>\n",
       "      <td id=\"T_b3b1d_row0_col0\" class=\"data row0 col0\" >2099</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row1\" class=\"row_heading level0 row1\" >hotpotqa</th>\n",
       "      <td id=\"T_b3b1d_row1_col0\" class=\"data row1 col0\" >1630</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row2\" class=\"row_heading level0 row2\" >strategyqa</th>\n",
       "      <td id=\"T_b3b1d_row2_col0\" class=\"data row2 col0\" >1595</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row3\" class=\"row_heading level0 row3\" >wiki_table_questions</th>\n",
       "      <td id=\"T_b3b1d_row3_col0\" class=\"data row3 col0\" >1129</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row4\" class=\"row_heading level0 row4\" >APPS</th>\n",
       "      <td id=\"T_b3b1d_row4_col0\" class=\"data row4 col0\" >667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row5\" class=\"row_heading level0 row5\" >algebra</th>\n",
       "      <td id=\"T_b3b1d_row5_col0\" class=\"data row5 col0\" >467</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row6\" class=\"row_heading level0 row6\" >number_theory</th>\n",
       "      <td id=\"T_b3b1d_row6_col0\" class=\"data row6 col0\" >270</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row7\" class=\"row_heading level0 row7\" >intermediate_algebra</th>\n",
       "      <td id=\"T_b3b1d_row7_col0\" class=\"data row7 col0\" >236</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row8\" class=\"row_heading level0 row8\" >prealgebra</th>\n",
       "      <td id=\"T_b3b1d_row8_col0\" class=\"data row8 col0\" >234</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row9\" class=\"row_heading level0 row9\" >counting_and_probability</th>\n",
       "      <td id=\"T_b3b1d_row9_col0\" class=\"data row9 col0\" >120</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row10\" class=\"row_heading level0 row10\" >precalculus</th>\n",
       "      <td id=\"T_b3b1d_row10_col0\" class=\"data row10 col0\" >73</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_b3b1d_level0_row11\" class=\"row_heading level0 row11\" >geometry</th>\n",
       "      <td id=\"T_b3b1d_row11_col0\" class=\"data row11 col0\" >60</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ffa111fa430>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "medium_instances = completed_results_w_stats[~is_failed].query(\"n_turns > 2\")\n",
    "print(f\"Number of medium instances: {len(medium_instances)} / {len(completed_results_w_stats)} ({len(medium_instances) / len(completed_results_w_stats):.2%})\")\n",
    "\n",
    "medium_instances_task_name_id_pairs = set(medium_instances[[\"task_type\", \"task_name\", \"task_id\"]].apply(tuple, axis=1))\n",
    "print(f\"Number of unique medium instances: {len(medium_instances_task_name_id_pairs)} / {len(all_unique_instances_task_name_id_pairs)} ({len(medium_instances_task_name_id_pairs) / len(all_unique_instances_task_name_id_pairs):.2%})\")\n",
    "\n",
    "medium_instances\\\n",
    "    .groupby(\"task_name\")[\"success\"].count().rename(\"n_tasks\")\\\n",
    "    .sort_values(ascending=False).to_frame().style.background_gradient(cmap=\"Blues\", axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 667 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/apps.jsonl\n",
      "Saved 2099 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/alfworld.jsonl\n",
      "Saved 467 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/algebra.jsonl\n",
      "Saved 120 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/counting_and_probability.jsonl\n",
      "Saved 60 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/geometry.jsonl\n",
      "Saved 1630 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/hotpotqa.jsonl\n",
      "Saved 236 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/intermediate_algebra.jsonl\n",
      "Saved 270 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/number_theory.jsonl\n",
      "Saved 234 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/prealgebra.jsonl\n",
      "Saved 73 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/precalculus.jsonl\n",
      "Saved 1595 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/strategyqa.jsonl\n",
      "Saved 1129 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/medium_instances/wiki_table_questions.jsonl\n"
     ]
    }
   ],
   "source": [
    "save_instances(medium_instances, os.path.join(ROOT_DIR, \"data\", \"trajectories\", \"medium_instances\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hard Tier - Need to involve GPT-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of unsolvable instances: 7635 / 21868 (34.91%)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>task_type</th>\n",
       "      <th>task_name</th>\n",
       "      <th>task_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>reasoning</td>\n",
       "      <td>hotpotqa</td>\n",
       "      <td>9900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>reasoning</td>\n",
       "      <td>prealgebra</td>\n",
       "      <td>875</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>code_generation</td>\n",
       "      <td>APPS</td>\n",
       "      <td>2558</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>code_generation</td>\n",
       "      <td>APPS</td>\n",
       "      <td>4544</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>decision_making</td>\n",
       "      <td>alfworld</td>\n",
       "      <td>look_at_obj_in_light-Book-None-DeskLamp-305/tr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7630</th>\n",
       "      <td>reasoning</td>\n",
       "      <td>hotpotqa</td>\n",
       "      <td>15148</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7631</th>\n",
       "      <td>code_generation</td>\n",
       "      <td>APPS</td>\n",
       "      <td>3605</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7632</th>\n",
       "      <td>reasoning</td>\n",
       "      <td>precalculus</td>\n",
       "      <td>5087</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7633</th>\n",
       "      <td>code_generation</td>\n",
       "      <td>APPS</td>\n",
       "      <td>1390</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7634</th>\n",
       "      <td>decision_making</td>\n",
       "      <td>alfworld</td>\n",
       "      <td>pick_two_obj_and_place-Pencil-None-Shelf-301/t...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>7635 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            task_type    task_name  \\\n",
       "0           reasoning     hotpotqa   \n",
       "1           reasoning   prealgebra   \n",
       "2     code_generation         APPS   \n",
       "3     code_generation         APPS   \n",
       "4     decision_making     alfworld   \n",
       "...               ...          ...   \n",
       "7630        reasoning     hotpotqa   \n",
       "7631  code_generation         APPS   \n",
       "7632        reasoning  precalculus   \n",
       "7633  code_generation         APPS   \n",
       "7634  decision_making     alfworld   \n",
       "\n",
       "                                                task_id  \n",
       "0                                                  9900  \n",
       "1                                                   875  \n",
       "2                                                  2558  \n",
       "3                                                  4544  \n",
       "4     look_at_obj_in_light-Book-None-DeskLamp-305/tr...  \n",
       "...                                                 ...  \n",
       "7630                                              15148  \n",
       "7631                                               3605  \n",
       "7632                                               5087  \n",
       "7633                                               1390  \n",
       "7634  pick_two_obj_and_place-Pencil-None-Shelf-301/t...  \n",
       "\n",
       "[7635 rows x 3 columns]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unsolvable_name_id_pairs = all_unique_instances_task_name_id_pairs - easy_instances_task_name_id_pairs - medium_instances_task_name_id_pairs\n",
    "print(f\"Number of unsolvable instances: {len(unsolvable_name_id_pairs)} / {len(all_unique_instances_task_name_id_pairs)} ({len(unsolvable_name_id_pairs) / len(all_unique_instances_task_name_id_pairs):.2%})\")\n",
    "unsolvable_df = pd.DataFrame(list(unsolvable_name_id_pairs), columns=[\"task_type\", \"task_name\", \"task_id\"])\n",
    "unsolvable_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "reasoning          3836\n",
       "code_generation    2003\n",
       "tabular            1010\n",
       "decision_making     786\n",
       "Name: task_type, dtype: int64"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unsolvable_df[\"task_type\"].value_counts()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task type: code_generation, task name: APPS, # of task_ids: 2003\n",
      "Saving 2003 instances (original: 4439) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/code_generation/train/apps.jsonl\n",
      "Task type: decision_making, task name: alfworld, # of task_ids: 786\n",
      "Saving 786 ALFWorld hard instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/decision_making/train/alfworld.txt\n",
      "Task type: reasoning, task name: algebra, # of task_ids: 334\n",
      "Saving 334 instances (original: 1226) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/algebra.jsonl\n",
      "Task type: reasoning, task name: counting_and_probability, # of task_ids: 311\n",
      "Saving 311 instances (original: 602) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/counting_and_probability.jsonl\n",
      "Task type: reasoning, task name: geometry, # of task_ids: 545\n",
      "Saving 545 instances (original: 727) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/geometry.jsonl\n",
      "Task type: reasoning, task name: hotpotqa, # of task_ids: 1108\n",
      "Saving 1108 instances (original: 3000) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/hotpotqa.jsonl\n",
      "Task type: reasoning, task name: intermediate_algebra, # of task_ids: 704\n",
      "Saving 704 instances (original: 1070) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/intermediate_algebra.jsonl\n",
      "Task type: reasoning, task name: number_theory, # of task_ids: 252\n",
      "Saving 252 instances (original: 691) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/number_theory.jsonl\n",
      "Task type: reasoning, task name: prealgebra, # of task_ids: 172\n",
      "Saving 172 instances (original: 750) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/prealgebra.jsonl\n",
      "Task type: reasoning, task name: precalculus, # of task_ids: 381\n",
      "Saving 381 instances (original: 520) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/precalculus.jsonl\n",
      "Task type: reasoning, task name: strategyqa, # of task_ids: 29\n",
      "Saving 29 instances (original: 2290) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/reasoning/train/strategyqa.jsonl\n",
      "Task type: tabular, task name: wiki_table_questions, # of task_ids: 1010\n",
      "Saving 1010 instances (original: 3000) to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_instances/raw/tabular/train/wiki_table_questions.jsonl\n"
     ]
    }
   ],
   "source": [
    "# Save Raw file for unsolvable instances (for GPT-4 generation)\n",
    "\n",
    "output_dir = os.path.join(ROOT_DIR, \"data\", \"trajectories\", \"hard_instances\", \"raw\")\n",
    "for (task_type, task_name), cur_df in unsolvable_df.groupby([\"task_type\", \"task_name\"]):\n",
    "    # get all the task_ids and select those from the original dataset into new df\n",
    "    task_ids = set(cur_df[\"task_id\"].unique())\n",
    "    print(f\"Task type: {task_type}, task name: {task_name}, # of task_ids: {len(task_ids)}\")\n",
    "\n",
    "    if not task_name == \"alfworld\":        \n",
    "        original_data_path = os.path.join(ROOT_DIR, \"data\", \"processed\", task_type, \"train\", f\"{task_name.lower()}.jsonl\")\n",
    "        original_df = pd.read_json(original_data_path, orient=\"records\", lines=True)\n",
    "        \n",
    "        # filter by task_ids\n",
    "        filtered_df = original_df[original_df[\"id\"].isin(task_ids)]\n",
    "\n",
    "        # save to output dir\n",
    "        cur_output_filepath = os.path.join(output_dir, task_type, \"train\", f\"{task_name.lower()}.jsonl\")\n",
    "        pathlib.Path(os.path.dirname(cur_output_filepath)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "        print(f\"Saving {len(filtered_df)} instances (original: {len(original_df)}) to {cur_output_filepath}\")\n",
    "        filtered_df.to_json(\n",
    "            cur_output_filepath,\n",
    "            orient=\"records\",\n",
    "            lines=True  \n",
    "        )\n",
    "    else:\n",
    "        # Only save the task_ids to a text file\n",
    "        cur_output_filepath = os.path.join(output_dir, task_type, \"train\", f\"{task_name}.txt\")\n",
    "        pathlib.Path(os.path.dirname(cur_output_filepath)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "        print(f\"Saving {len(task_ids)} ALFWorld hard instances to {cur_output_filepath}\")\n",
    "\n",
    "        with open(cur_output_filepath, \"w\") as f:\n",
    "            for task_id in task_ids:\n",
    "                f.write(f\"{task_id}\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPT-4 Solved Instances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of medium instances: 423 / 3574 (11.84%)\n",
      "Number of unique medium instances: 423 / 3574 (11.84%)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_63152_row0_col0 {\n",
       "  background-color: #08306b;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_63152_row1_col0 {\n",
       "  background-color: #1561a9;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_63152_row2_col0 {\n",
       "  background-color: #2272b6;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_63152_row3_col0 {\n",
       "  background-color: #58a1cf;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_63152_row4_col0 {\n",
       "  background-color: #5da5d1;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_63152_row5_col0 {\n",
       "  background-color: #7cb7da;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_63152_row6_col0 {\n",
       "  background-color: #a1cbe2;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_63152_row7_col0 {\n",
       "  background-color: #b0d2e7;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_63152_row8_col0 {\n",
       "  background-color: #d5e5f4;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_63152_row9_col0 {\n",
       "  background-color: #f7fbff;\n",
       "  color: #000000;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_63152\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_63152_level0_col0\" class=\"col_heading level0 col0\" >n_tasks</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >task_name</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row0\" class=\"row_heading level0 row0\" >APPS</th>\n",
       "      <td id=\"T_63152_row0_col0\" class=\"data row0 col0\" >74</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row1\" class=\"row_heading level0 row1\" >algebra</th>\n",
       "      <td id=\"T_63152_row1_col0\" class=\"data row1 col0\" >62</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row2\" class=\"row_heading level0 row2\" >intermediate_algebra</th>\n",
       "      <td id=\"T_63152_row2_col0\" class=\"data row2 col0\" >58</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row3\" class=\"row_heading level0 row3\" >hotpotqa</th>\n",
       "      <td id=\"T_63152_row3_col0\" class=\"data row3 col0\" >46</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row4\" class=\"row_heading level0 row4\" >number_theory</th>\n",
       "      <td id=\"T_63152_row4_col0\" class=\"data row4 col0\" >45</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row5\" class=\"row_heading level0 row5\" >geometry</th>\n",
       "      <td id=\"T_63152_row5_col0\" class=\"data row5 col0\" >40</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row6\" class=\"row_heading level0 row6\" >counting_and_probability</th>\n",
       "      <td id=\"T_63152_row6_col0\" class=\"data row6 col0\" >34</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row7\" class=\"row_heading level0 row7\" >precalculus</th>\n",
       "      <td id=\"T_63152_row7_col0\" class=\"data row7 col0\" >31</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row8\" class=\"row_heading level0 row8\" >prealgebra</th>\n",
       "      <td id=\"T_63152_row8_col0\" class=\"data row8 col0\" >22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_63152_level0_row9\" class=\"row_heading level0 row9\" >strategyqa</th>\n",
       "      <td id=\"T_63152_row9_col0\" class=\"data row9 col0\" >11</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7ff9f1684940>"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gpt4_is_failed = gpt4_completed_results_w_stats[\"success\"].apply(lambda x: not bool(x))\n",
    "gpt4_instances = gpt4_completed_results_w_stats[~gpt4_is_failed].query(\"n_turns > 2\")\n",
    "print(f\"Number of medium instances: {len(gpt4_instances)} / {len(gpt4_completed_results_w_stats)} ({len(gpt4_instances) / len(gpt4_completed_results_w_stats):.2%})\")\n",
    "\n",
    "gpt4_all_unique_instances_task_name_id_pairs = set(gpt4_completed_results_w_stats[[\"task_type\", \"task_name\", \"task_id\"]].apply(tuple, axis=1))\n",
    "\n",
    "gpt4_instances_task_name_id_pairs = set(gpt4_instances[[\"task_type\", \"task_name\", \"task_id\"]].apply(tuple, axis=1))\n",
    "print(f\"Number of unique medium instances: {len(gpt4_instances_task_name_id_pairs)} / {len(gpt4_all_unique_instances_task_name_id_pairs)} ({len(gpt4_instances_task_name_id_pairs) / len(gpt4_all_unique_instances_task_name_id_pairs):.2%})\")\n",
    "\n",
    "gpt4_instances\\\n",
    "    .groupby(\"task_name\")[\"success\"].count().rename(\"n_tasks\")\\\n",
    "    .sort_values(ascending=False).to_frame().style.background_gradient(cmap=\"Blues\", axis=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 74 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/apps.jsonl\n",
      "Saved 62 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/algebra.jsonl\n",
      "Saved 34 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/counting_and_probability.jsonl\n",
      "Saved 40 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/geometry.jsonl\n",
      "Saved 46 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/hotpotqa.jsonl\n",
      "Saved 58 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/intermediate_algebra.jsonl\n",
      "Saved 45 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/number_theory.jsonl\n",
      "Saved 22 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/prealgebra.jsonl\n",
      "Saved 31 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/precalculus.jsonl\n",
      "Saved 11 instances to /shared/nas2/xingyao6/projects/llm-agent/data/trajectories/hard_gpt4_instances/strategyqa.jsonl\n"
     ]
    }
   ],
   "source": [
    "save_instances(gpt4_instances, os.path.join(ROOT_DIR, \"data\", \"trajectories\", \"hard_gpt4_instances\"))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "interaction-eval",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
